Rao-Blackwellised Reparameterisation Gradients
Authors: Kevin H. Lam
Thang D Bui
George Deligiannidis
Yee Whye Teh
Venue: NeurIPS 2025 poster ()
Keywords: reparameterisation trick
stochastic gradient estimation
Rao-Blackwellisation
variance reduction
Primary area: probabilistic_methods
TL;DR: We present the Rao-Blackwellised reparameterisation gradient estimator and show its performance gains on a suite of probabilistic models.

Abstract: Latent Gaussian variables have been popularised in probabilistic machine learning. In turn, gradient estimators are the machinery that facilitates gradient-based optimisation for models with latent Gaussian variables. The reparameterisation trick is often used as the default estimator as it is simple to implement and yields low-variance gradients for variational inference. In this work, we propose the R2-G2 estimator as the Rao-Blackwellisation of the reparameterisation gradient estimator. Interestingly, we show that the local reparameterisation gradient estimator for Bayesian MLPs is an instance of the R2-G2 estimator and Rao-Blackwellisation. This lets us extend benefits of Rao-Blackwellised gradients to a suite of probabilistic models. We show that initial training with R2-G2 consistently yields better performance in models with multiple applications of the reparameterisation trick.

Rao-Blackwellised Reparameterisation Gradients
Kevin H. Lam Department of Statistics University of Oxford
Thang D. Bui School Computing Australian National University
George Deligiannidis Department of Statistics University of Oxford
Yee Whye Teh Department Statistics University of Oxford
(lam,deligian) www..//..t..ettttts.oomm ac.xa, thang. bui@anu. edu.
Abstract
Latent Gaussian variables have been popularised in probabilistic machine learning. In turn, gradient estimators are the machinery that facilitates gradient-based opti- misation for models with latent Gaussian variables. The reparameterisation trick is often used as the default estimator as it simple implement and yields low- variance gradients for variational inference. In this work, we propose the R2-G2 estimator as the Rao-Bl Rao-Blackwellissee Interestingly, we show that the local reparameterisation gradient estimator for Bayesian MLPs an instance of the R2-G2 estimator and Rao-Bi Rao-Blackwe Blackwellisation. This lets us extend benefits of Rao- Soo-Blackwellised gradients to suite of probabilis- tie models. We show that initial training with R2-G2 consistently yields better performance in models with multiple applications of the reparameterisation trick.
the reparameterisation gradient estimator.
1
Introduction
Latent random variables are ubiquitous in probabilistic machine learning (ML) they enable to embed prior assumptions into models, such as uncertainty in model parameters and low-dimensional latent representations of observed data structures. Thus, latent random variables appear range of modelling tasks, including variational inference, generative modelling, and density estimation. A wide variety of modern models with latent random variables, particularly those that parameterise a neural network (NN), are trained with gradient- based optimisation. This presents non-trivial problem on how analytical derivatives can incorporate stochasticity, and has led to substantial interest gradient estimators. These estimators enable gradient-basee based optimisation by computing approximations of gradients that are compatible with automatic differentiation software. For comprehensive reviews of gradient estimators with latent continuous and discrete random variables, refer the reader to Mohamed et al. (2020) and Huijben et al. (2023) respectively.
We revisit the class of gradient estimators for latent Gaussian variables (Kingma and Welling 2014 Rezende et al., 2014 Titsias and aGeeill 2014), Gaussian variables often serve as the default distribution for the noise prior in probabilistic machine learning tasks. These estimators, known as pathwise estimators the reparameterisation trick, have been popularised as they produce low-variance gradients for variational inference. In particular, it has been shown in Xu et al. (2019) that reparameterisation gradient estimator has lower variance than Bao-Blackwellised REINFORCE estimator (Williams) (992)) for variational inference. These estimators are also attractive as they are single-sample estimators since training models with them only requires sampling operations be done
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

once. In turn, this also limits costly function evaluations needed for training models grow larger in the era of deep learning.
In this work, we present the R2-G2 estimator as an extension of the reparameterisation gradient estimator formed by conditioning on pre-activations in backpropagation. Notably, show that the local reparameterisation gradient estimator proposed in Kingma et al. (2015) is an instance of the R2-G2 estimator for linear layers in a Bayesian MLP with independent weights and establish that is equivalent to Bho-Blackwelliscd reparameterisation gradient estimator.
Our main contributions are follows:
We present R2-G2, novel, general-purpose and single-sample gradient estimator for latent Gaussian variables as Rao-Bl Blackwellised reparameterisation gradient estimator;
We show that the local reparameterisation gradient estimator is an instance of the R2-G2 estimator and enjoys variance reduction in gradients from Rao- Blackwellisaaionn
We empirically show that initial training with the R2-G2 estimator consistently yields higher likelihoods on Bayesian Neural Networks and higher ELBOs on hierarchical Variational Autoencoders than the reparameterisation gradient estimator.
2 Problem setting
Let v E R be a vector continuous random variables where we have v ~ go for some probability provided continuously differentiable
distribution 9 parameterised by vector 0. Suppose loss function f R" > R that depends on the random variables v. To enable gradient- based training of parametric model f, our goal to compute the gradient of the expected loss
Vg [v)]
(1)
where 'D is dataset and f' is parametric function optimised by the loss such as neural network (NN), with both affecting the evaluation of é (i.e. lD.r). We note that described generally here to encompass both supervised and unsupervised problems. For case of notation in the remainder of this work, we will shorten our notation the loss from fD,((v, 0) to where the context is clear with its dependence v, 0, D, and implied. Popular choices of l aim to maximise the log likkliihoodd or the evidence lower bound (ELBO) for the dataset D. The former involves setting as the negative log- ligeliie LNLL. The latter sets l as
20 (v (D) p(v)
(ELBO(V) = (DLv( + log
where p is a prior distribution of v and setting ge the posterior of its approximation, making Equation ! equivalent to maximising the ELBO. Moreover, is equivalent to variational Bayesian inference when the form of go is restricted as it becomes approximation of the posterior.
In practice, Monte Carlo approximation Egg [e(v,
M EM ((v(), 0) to write
E [[(v)]
Vol(v(1)
M
where v1), ..... ((f) 90- However, this would require M evaluations of the loss & which is not desirable when the underlying model f is large or costly to evaluate. In this work, we focus on the setting where independent Gaussian random variables with the aim to derive single-sample gradient estimator (i.e. M = 1) that is unbiased and enjoys reduced variance.
Related work
In this section, we revisit the single-sample gradient estimators compatible with Gaussian random variables, namely REINFORCE, reparameterisatiio trick and local reparameterisatioo trick.

---

REINFORCE The REINFORCE gradient estimator, also known as the score function estimator, is proposed by using the der-deivativee trick:
Vo log go (v)
Voge(v) which comes result of the se(v)
chain rule Glynn 1990 Williams
1992). Formally, the partial derivatives of parameters 0 are given
by the REINFORCE estimator
SCORE
=-- (Vo log qo)
It is an unbiased estimator of Vo Eng [e]. and only requires that one is able to evaluate and sample from it. The latter requirement is not restrictive and easily achieved in most settings since the form of de is often assumed as part of model specification such as variational inference. Despite these desirable properties, it is well-known that the REINFORCE estimator suffers from high variance (Greensmith et al. 2001 Xu et al.) 2019).
Reparameterisation Trick The reparameterisation trick is compatible with random variables that have a location-scale parameterisation or have tractable inverse cumulative distribution functions (CDFs) 1958, Glasserman, 2003] Kingma and Welling 2014 Rezende al. 2014 Titsias
and Gredill Figurnov et al., [0188] (0018)Jana Jankowink and Obermeyer, 2018). 8). For Gaussian
random variables, they are reparameterissale using a location-scall transformation which lets us write
V N g M, V)= 4 Ve
where V E RRX" Σ VVT and (o(e) = VV I). The latter expression for v is compatible with automatic differentiation and enables gradient- based updates to μ and V. Typically, V is diagonal matrix parameterised by vector of variances = (of.....0% 2.) That is, V = Σ with = diag(T), so we have = and apply the chain rule to yield the reparameterisation gradient estimator
{u, T} as our training parameters. Formally, we can write ((g(e,0)
vol RT
(Je(v) I,
3
(1,ET) I)
(2)
where Je is the Jacobian of with respect to v, and [A[B] denotes partitioned matrix with block matrices A and B. The terms in the matrix are derived from the Jacobians 3,(44) = I, and (T) A @ (1,ET). Equation 2 is unbiased estimator of VeE., [e] due to the equivalence in distribution achieved by the reparameterisation trick. In the context of probabilistic modelling, the reparameterisation trick is often applied to individual scalar inputs of decoder within variational autoencoder (VAE) or individual weights Bayesian neural network (BNN). The latter may also be referred as the global reparameterisation trick.
Local Reparameterisation Trick The local reparamcterisation trick was proposed for NNs where the weights of linear layers independent Gaussian random variables (Kingma et al.) 2015). Given an input x E IR",, the pre-activations of these linear layers are given by
x g(e (())
z
E R
x g(e(m)) (((m))
where ( and e(i) are independent for all 4 j. and they admit a factorised Gaussian distribution in = 11 2z, where each pre-activation z; dz.. Instead of sampling & and computing x g, is more efficient to locally apply the reparameterisation trick to directly sample scalar pre-activations
24 N
Σ 13
Σ j=1
(1))
x (0,00
Σ
2
E
(3)
1=1
where y and (6) = (0 (1)
( are the global mean and variance parameters of the Gaussian variables used to compute 24 respectively, (()) (,()) (())
]. and & ~ 9o (0, 1). The latter expression in Equation is known as the local reparameterisation trick as it changes the parameterisation of each 2; scalar function of (4) and g(s), Formally, we can write E(g(e, 0))
é (z1 (1)) ..... Zm (Sm, 0 ((m))) where 8 p E denotes functions applied to to compute
3

---

Table 1: Examples of linear maps applied to Gaussian variables in probabilistic neural networks.
Application
W
Ve
BNN
Outputs of previous hidden layer
Gaussian weights in a linear layer
BNN
Vectorised patches of an image
Gaussian weights in a convolutional layer Gaussian latent variables
VAE
Linear layer after reparameterisation
the loss (i.e. non-linearities and upper NN layers), and apply the chain rule to yield the local reparameterisation gradient estimator for each 0())
LRT Vool
ae
1 1 (o & (xxx) 1)
T
V
x
0z:
where @ denotes element-wise product of matrices. The terms in the matrix are derived from
the Jacobians Jz,( (x())) = =x and (z,(7()) = ( (of))
& (xxx) Equation is an unbiased estimator of Vol [8] due to the equivalence in distribution achieved by the local reparameterisation trick. A notable application of the local reparameterisation trick is when performing mean-field variational inference within the linear layer of BNN. It has been empirically observed that Equation 4 has lower variance than Equation 2 (Kingma et al. 2015). In Theorem
4.3 show that the local reparameterisation gradient estimator equivalent applying Rao- Blackwellisation to the reparameterisation gradient estimator and thereby enjoys variance reduction benefits of Bloellaaccellssooonnn
4 R2-G2 Gradient Estimator
The reparameterisation trick can be seen as procedure that couples the specification of a probabilistic model with an expression for gradients. turn, assumptions on former would restrict the applicability of the latter. This highlights the disadvantage of the local reparameterisation trick: the probabilistic model and gradients induced by the local reparameterisation trick is not suitable for general settings where the pre- ae-aaiivvtions do not admit a factorised Gaussian distribution (i.c. the covariance matrix of z is not diagonal). On the other hand, the local reparameterisation trick has heen empirically shown enjoy lower variance of gradients than the global reparameterisation trick.
To get the best of both worlds, we seek reparameterisation gradient estimator that is and enjoys reduced variance. To this end, present our contribution: the Bao-Blackwellisse Reparameterisation Gradient Estimator for Gaussian random variables, coined the R2-G2 estimator, as the realisation of multivariate Gaussian vectors resultant from linear transformations of Gaussian vectors. of the reparameterisation gradient estimator. We then
of the reparameterisation gradient estimator by conditioning the
We first describe the Rao- Rao-Blackw Blackwellisatinn provide the analytical form of the R2-G2 estimator and a summary of key properties the R2-G2 estimator and its connection to existing methods. In particular, show that the local reparameteri- sation gradient estimator is an instance the R2-G2 estimator. We then conclude with practical implementation bypass costly matrix inversion operations by reformulating matrix-vector products as the solution to a quadratic optimisation problem. We defer all proofs to Appendix B
4.1 BaollBaacwelllisation of the Reparameterisation Gradient Estimator
The idea behind Rao-Blackwer nho-Blackwelllsattoon scheme is condition linear transformations of Gaussian variables since conditional Gaussian distributions can be deseribed analytically. This is further motivated by the observation that
decompose loss evaluations in NNs
(5) where W R - Rm is a linear map and W g(e f) are pre-activations. This generalises the idea of the local reparameterisation trick where W is row vector (i.e. W = = and inputs of are m scalar pre-activations. In deep learning, W frequently appears hidden layers within NNs. A list of
E(g(E, 0)) = (ē ()(g(e

---

common linear transformations in probabilisticc NNs is given in Table Applying the chain rule to gives an alternative expression of the reparameterisation gradient estimator
Equation
V WJ g) W I, ト (1,Ee) ])
(6)
where J/ is the Jacobian of É. With Equation 6 hand, we now present the R2-G2 estimato
Definition 4.1 (R2-G2) The R2-G2 gradient estimator is given by
Vel = =Eio
RT Vol(c)
J&(W 9) W. Eno In
E o (IneT)
(7)
where W 9 = ~ go with mean E" = A (AAT) (z Wp) A (AAT)) Ae with A WV. By linearity of expectations, also admits closed-form expression
R2-G2
[Ja(W 9) W [I, E ((o()()))
V
(8)
The R2-G2 estimator is a single-sample gradient estimator as it is a conditional expectation of the reparamcterisation gradient estimator formed by conditioning pre-activations z, making it the Rao-Blackwel Blackwellisation (Blackwell 1947; Rao et al. 1992) of the reparameterisation gradient estimator from Equation 6 The law of total variance lets deduce that the R2-G2 estimator has lower variance since it swaps the random matrix in Equation with its conditional expectation. We conclude description of the R2-G2 estimator with its key properties, namely unbiasedness and enjoying lower variance than the reparameterisation gradient estimator.
single sample of the vector of
Proposition 4.2 Denote Oz N (W u1, AA
Then we have
R2-G2
R
V
V R
and
R2-G2 Vol
Vol RT
Eqg
Vo
V
Ego
Vol
V
[9]
4.2
Computation of nho-Blackwellisatioon scheme least-squares problem
The R2-G2 estimator exploits the analytical form of the conditional Gaussian distribution. This done by either sampling from the conditional Gaussian distribution go within Equation or directly using the mean of the conditional Gaussian distribution e" in place of within Equation 8 Equation 7 requires computing the Cholesky factor of the covariance matrix of go. which incurs (m)²) computational cost and O(m²) storage cost. The latter makes impractical for training NNs as
each gradient descent step would require instantiating the matrix AA T This makes Equation g more desirable as long as e" is computed and stored efficiently. To do so, we consider the linear system
AA B Ae.
It can be verified that any solution 8 to Equation 9 satisfies e" A ATB (see Appendix D
Least-squares characterisation of Rao- www.-Blkkelllsttiiiinn Observe that Equation 9 is normal
equation that describes the first-order optimality condition of the quadratic optimisation problem
exp -- A ATBI3
minimise BERm
Glle
A B
< maximise BERm
We can interpret computing solution B as fitting single-sample multivariate linear model e/A
A A
'Equations 7 and g in Definition 4.1 only hold when are independent Gaussian variables. A more general expression be derived when independence does not hold.

---

where B E R and ò ~ (O,, I.). Under this linear model, the goal is to solve for 8 by using covariance parameters as predictors of the noise e. With 8 computed. we can then compute
E A A (AAT)* Ae
(10)
which corresponds the expression of the mean of a conditional Gaussian distribution (see Lemma A.I.. In other words, computing the R2-G2 estimator is inherently solving a least squares problem and applying a linear transformation. Intuitively, the former the mechanism giving variance reduction. In the language of linear models, the observed noise is projected to the fitted noise e" which has minimal Euclidean norm.
Iterative solver By defining functions that compute matrix-vector products with the ma- trix AA we can use the conjugate gradient algorithm to solve Equation 9 for 8 The conjugate gradient algorithm is an iterative method that terminates
Algorithm Forward Pass with R2-G2 Gradients.
Input: matrix A, noise vector e. Compute Ae.
Compute 8 = con conjung jugate_gradiont(A 2).
Compute E" A TT"
most rank () S
Compute z Ac'.
m iterations (Kaasschieter] 1988 Nocedal and Wright T999| HHaamil 2018), and only re-
Output: stop-gras ceop-gradienteo (2-
quires storing the updated solution at each iteration. While the former implies that the worst-case computational costs remains at (mm)), the latter implies the storage cost is reduced to ((m) and makes it practical to compute B In our setting, the worst-case computational costs of O(m²) can be reduced since the structure of A is known (see Appendix F). With 3 in hand, we can compute e* and modify forward computations to use the R2-G2 gradient estimator for backpropagatior by implementing Algorithm |
in deep learning frameworks such as PyTorch (Paszke et al. 2019).
4.3 Connections to related work
= x and the conditional distribution e |2, ~ 40 for i 1,.... m within Definition 4.1 we are able to show that the local reparameterisation gradient estimator is equivalent to the R2-G2 estimator. presenet Theorem 4.3 presents this equivalence and formalises the empirical variance reduction of the local reparameterisation trick observed in the experiments of Kingmaa et al. (2015), as an instance of Rao- aaalllackwlllisationn and the R2-G2 estimator. Formally, it shows the local reparameterisation trick is equivalent to using pre-activation samples for forward computations and a Rao- Boo-BBaccwellseed global reparameterisation gradient estimator to update parameters in linear layers.
Local Reparameterisation Gradient Estimator
By setting ( = (diag(T
.(6)
Theorem 4.3 Suppose we have a BNN linear layer where weights are independent Gaussian random variables. That is, g(3) = {,l),~( (0)}
where 7(6)
(1)
(8) 2
for 1,.... m.
R2-G2
RT
LRT
Then for each 1,..., .,, m, we have Vocsl
V F l
Vocal
l
and
LIeT
17
F Rr
Eg
Vomel
V E
Rec [8]
S
E 90
V
EE 00))
[C]
As an analogy, we can view the way that R2-G2 generalises the local reparameterisation gradient estimator, in the way that square matrix inversion generalises scalar inversion. Recall that Equation 3 exploits the fact that the Gaussian pre-activations z of mean-field BNN linear layers have a diagonal covariance matrix by design. For forward computations, this means each can be efficiently sampled
with the reparameterisation trick by using the square root of the variance 02, -ド
The subtle benefit of reduced variance in gradients is due to the square root function having derivative which matches its reciprocal (up scaling factor). This means that the variance o² (i.e. scalar) is inverted when using the local reparameterisation gradient estimator. The R2-G2 estimator naturally extends the inversion of scalar variances to square covariance matrices, by calling the conjugate gradient algorithm, thereby making it suitable for other probabilistic models.
In Algorithm I computing zt ensures automatic differentiation calculates Equation 8 for backpropagation. ))+ + x² ensures the forward computation is still performed with z.
while the output stop_ grooppggrdeennt (z- z²)

---

Table 2: Log-likelihoods and classification accuracies (%) of BNNs using the R2-G2, Reparameteri- sation (RT) and Local Reparameterisation (LRT)) estimators over 5 runs. Higher better, Error bars denote 1.99 standard errors (a/ /55 over runs. See text for details.
Log-likelihood
Accuracy
Dataset
Estimator
Train
Test
Train
Test
MNIST
R2-G2
-3.00
3.06
99.81 = 0.03
98.00 = 0.08
LRT
-3.00
99.81 0.04 -3.06 0.00
97.99 + 0.07
RT
-.00
-3.07 + 0.00
09.85
98.00 0.11
CIFAR-10
R2-G2
3.83
-3.87 0.01
70.08 0.43
67.97 + 0.66
RT
-3.84 0.01
-3.88 + 0.02
69.76 = 0.46
67.98 + 0.88
Variance reduction Stochastic Linear Regression Variance reduction through stochastic linear regression has previously been explored in the context of variational inference Salimans and Knowles (2013) 2014); Salimans (2014). Specifically, variational inference is presented as fitting linear regression with the unnormalised log posterior as the dependent variable and the sufficient statistics of latent random variables as explanatory variables. The connection to variance reduction of gradients is then made by using multiple samples to construct control variates that correlate with the gradient of the Kl.-divergence. This differs from our work as we present variance reduction of gradients as Bto-Blackwellln explicate the connection to linear regression with the reparameterisation noise and covariance matrix parameters the dependent and explanatory variables respectively.
through the R2-G2 estimator as single-sample gradient estimator, and
5 Experiments
5.1
Protocol
The R2-G2 estimator can readily applied to any existing application of the reparameterisation trick for Gaussian variables. While this may appear restrictive, we note that Gaussian distributions often serve as default prior or noise distribution in probabilistic ML tasks. To motivate our experiments, we note that the R2-G2 and reparameterisation gradient estimators would yield the same model when a lurge number of gradient descent steps are taken, as the mean gradient of both estimators are the same. The focus of our experiments is during initial training, where only small number of gradient descent steps are taken. Examples of scenarios where this setting can be beneficial include pre-training to discover better initialisation for model, fine-tuning pre-trained model, when the amount of training is limited by a computation budget.
We evaluate the benefits of initial training with the R2-G2 estimator for probabilistic models. We consider two standard tasks with variational Bayesian models that utilise a mean- field approximation of the posterior: image classification with BNNs and generative modelling with hierarchical VAEs. our experiments, we provide numerical comparisons against the reparameterisation
(RT) and local reparamcterisation (LRT) gradient estimators, where permissible. We do not compare against the REINFORCE estimator as is well-known that its high variance makes optimisation difficult. Unless stated otherwise, we compute pre-activations in the same way as the reparameterisation trick (i.e. sampling g(e, 0) and computing W g). We defer the full details of our experiments to Appendix B
5.2 Image classification with Bayesian Neural Networks
We consider fully stochastic BNNs for image classification tasks on two standard benchmark datasets: MNIST (LeCun et al.) 2010)) and CIFAR-10 (Krizhevsky and Hinton 2009). In this setting, the input weights (i.e. global variables). To enable
data D is dataset and the latent variables training with minibatches, we use the stochastic approximation of the variational lower bound
logp(y(x)= > Egg (v|p) log p(v)
B
log 1x)
(o(v(D))
B
7

---

where N is the size of the dataset D and { {x((), y B use the standard train and test splits of both datasets. See Table@ classification accuracies and log-likelihoods reported on the train and test sets each dataset.
1-1 Ds is minibatch of size B < N. We
LRT
Log Gradient Variance
R2-G2
-0.5
-1.5
-2.00
-2.5
10
15
20
Epoch
Log average gradient variance v.s. epoch for the top layer of Bayesian MLP trained MNIST over runs. We compare the variance of gradients when training using the reparameterisation (RT), local reparameterisation (LRT) and R2-G2 estimators.
Figure 1:
For MNIST, we used the two-layer multi-layer perceptron (MLP) architecture from Srivastava et al. (2014), which consists of two hidden layers of 1024 ReLU units. Each estimator is applied to all linear layers. For experiments with the R2-G2 estimator, we compute pre-activations
the same way as the local reparameterisation trick (i.e. sampling pre-activations directly). Our results illustrate that accuracies and ligelioossss for the R2-G2 and LRT estimators match exactly. In Figures [ [andd 3] the variance of gradients for the R2-G2 and LRT estimators also coincide and are much lower than those the RT estimator. These observations empirically supports Theorem 4.3 the LRT estimator is instance of the R2-G2 estimator and Rao-Blackwelll Blackwellisation. In other words, the local reparameterisation trick equivalent to directly sampling pre-activations forward computations
and using Rao- Bao-Blackwellised RT estimator in backpropagation.
For CIFAR-10, used the VGG-11 architecture described in Simonyan and Zisserman (0115). where each estimator is applied to the last four convolutional layers. For the last four convolutional layers, we do not include the local reparameterisation trick as benchmark since sampling scalar pre-activations directly would ignore dependencies induced by sharing weights in a convolutional layer (i.e. additionally assumes AAT T is diagonal). The RT estimator applied the first four convolutional layers. To limit the effect of gradient variances from linear layers on gradient variances of convolutional layers, we applied the LRT estimator to linear layers.
Across both datasets, we found that initial training with the R2-G2 estimator yields higher log- likelihoods than the RT estimator while enjoying similar levels of accuracy. These results extend the benefits of training with Rao- Boo-Blaccwellieed gradients from Bayesian MLPs to Bayesian CNNs.
5.3 Generative modelling with Hierarchical Variational Autoencoders
We consider VAEs for generative modelling tasks three standard benchmark datasets: MNIST (LeCun et al. 2010), Omniglot (Lake et al. 2015) and Fashion-MNIST (Xiao et al. 2017). We do not Tocus comparisons of gradient variances due to architecture constraints of VAEs. Validating Proposition requires independent and identically distributed (i.i.d.) samples of z ~ Te which requires computing Cholesky factor in each iteration. This is computationally expensive and not representative of computations in VAEs. In practice, we compute z = W g(c, 0) where W the linear layer following reparameterisation, which does not enable sampling the full support of 7z.
We found that training one-layer VAEs with the R2-G2 estimator did not guarantee performance gains (see Appendix and surmise that reducing gradient variance hinders the learning of stable

---

Table 3:
Test variational lower bounds for hierarchical VAEs using the R2-G2 and Reparameterisation (RT) estimators. Higher is better. Error bars denote 1.99 standard errors (o/v5) over 5 runs. See text for details.
VAE Layers
Estimator
MNIST
Omniglot
Pashion-MNIST
R2-G2 RT
-106.85 + 5.00
1229.0 0.74
-240.23 0.64 -240.59 + 0.93
-107.64 + 8.46
-131.48 + 1.66
R2-G2
-102.45 = 3.73
1134.95 = 2.26
-240.15 = 0.86 -240.89 + 0.97
RT
-111.50 + 5.40
-136.12 2.70
representations in this setting as only training of the encoder is affected by the R2-G2 estimator, An investigation of this phenomena is beyond the scope of this paper and leave this for future work.
We present the results of our experiments for two-layer and three-layer VAEs (i.e. hierarchical VAEs). In this setting, the input data DD is an observation and the latent variables v is its latent representation (i.e. v are local variables). The objective is maximise the variational lower bound on the log-likelihood
K D(D, v(0)) ==1 4o (v())p))
1 K
logp(x)= > Buoiv( D) (V()) log
where 'D denotes input data, and (v(0) T'=1 are vectors of latent Gaussian random variables. For training, we use a single sample (K = 1), which is equivalent to variational inference. For testing, we use 5000 samples (K 5000), which is equivalent to importance weighted variational inference, providing a tighter bound on the log-likelihood (Burda et al. 2016). For each L-layer VAE, we used bottom-up variational posterior and top-down generative process with latent variables {v.)_..
L-1
((1v))/=11)> ((vxlx) [(v p(x, (v.)/a)) p(x II p(vi/ve+1 Im1
2
with latent spaces of 50 units and each conditional distribution is parameterised by a MLP with two hidden layers of 200 tanh units (see Burda al. (2016); Bauer and Mnih (202T)). We used factorised Bernoulli likelihood, and factorised Gaussian variational posterior and prior.
We applied the R2-G2 estimator as single layer within the decoder that applies reparameterisation and the linear transformation that follows it (i.c. W and V are decoder parameters). We used dynamic binarisation of all three datasets (Salakhutdinov and Murray 2008). We used the standard train and test splits for MNIST and Fashion-MNIST, and the train and test split from (Burda et al. 2016) for Omniglot. See Table 3 for test lower bounds the log- likelihood of all datasets after 100,000 steps.
Across all datasets and hierarchical VAEs, we found that the R2-G2 estimator consistently yielded higher test ELBBOs than the RT' estimator. For three-layer VAFs, we saw substantial gains of 9.05 and 1.17 nats on MNIST and Omniglot respectively (see Figures and [) These results extend benefits of initial training with Rao-Blackwellised gradients, from Bayesian NNs to Hierarchical VAEs.
6 Conclusion
We have presented the R2-G2 estimator as a novel, general-purpose and single-sample gradient estimator for latent Gaussian random variables the Baa-llackwellisationn of the reparamcterisation gradient estimator. Our method motivated by the widespread usage of Gaussian distributions in probabilistic ML and applications of Rao- noo-Blackwellisationn in gradient estimators for latent discrete
variables (Liu et al. 2019 Dong et al.) 2020 Kool et al. 2020 Paulus et al. 2021).
We theoretically and empirically show that the local reparameterisatioo trick is an instance of Rao- Blackwellisation and the R2-G2 estimator for linear layers BNNs with independent weights. is equivalent to sampling pre-activations in forward computations and updating parameters with Rao- Bao-Blackwellised reparameterisation gradient estimator backpropagation. This explicates the empirical evidence that the local reparameterisation trick reduces variance of gradients obtained

---

-110
-110
ELBO
ELBO
-130
-130
-150
-150
20K
40K
60K
80K
100K
20K
40K
60K
80K
100K
Steps
Steps
RT: - R2-G2
Figure 2: Bounds log-likelihood v.s. optimisation steps for three-layer VAE trained MNIST over runs. We compare the bounds on log-likelihoods when training using the reparameterisation (RT) and R2-G2 estimators. Training with the R2-G2 estimator improves bounds log- le-likeli on both the training set (left) and test set (right).
by the global reparameterisation trick as the benefit of Rao- Blackwellisation stated by Theorem 4.3 By casting the local reparameterisation gradient estimator as Rao-Blackwellised gradients for Bayesian MIPs, we showed that initial training with Rao-Blackwellised gradients also yield gains in performance for other models such as Bayesian CNNs and hierarchical VAFs. While performance gains were sometimes modest, they were consistent across models with multiple applications of the reparameterisation trick and particularly prominent in hierarchical VAEs. The main limitation of the R2-G2 estimator is the computational cost from solving Equation 9 and its implicit dependence on how pre-activations the former, we exploited low dimensionality in upper convolutional layers of Bayesian CNNs and latent spaces of VAEs, thereby requiring less iterations of the conjugate gradient algorithm.
computed. The latter is inherited from model parameterisation. To mitigate
While our work has focused on the case where the matrix V diagonal, we note that R2-G2 can be extended to the non-diagonal case, such as low-rank covariance matrices (Tomczak et al., 2020), by defining appropriate matrix-vector product functions called within the conjugate gradient algorithm. Aside from likelihoods and ELBOSs, we note that the R2-G2 estimator can be applied to other objective functions. A potential extension of our work can to augment existing gradient estimators for variational inference in the multi-sample setting (Roeder et al. 2017 Rainforth al. [018 Tucker ctt al. 2019), Bauer and Mnih, 202T) with the R2-G2 estimator and evaluate its effectiveness. We leave these directions for future work.
7 Acknowledgements
We would like to thank Andriy Mnih for helpful discussions and comments on drafts this paper. Kevin H. Lam gratefully acknowledges his PhD funding from Google DeepMind. Thang D. Bui acknowledges support from the National Computing Infrastructure (NCI) Australia. George Deligian- nidis acknowledges support from the Engineering and Physical Sciences Research Council [grant number EP/Y018273/1). Yee Whye Teh acknowledges support from the Ministry of Digital Devel- opment and Information (MDDI) under the Singapore Global AI Visiting Professorship Program (Award No. 2024-002..
References
Matthias Bauer and Andriy Mnih. Generalized doubly reparameterized gradient estimators. In Pre- ceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 738-747, PMLR, 18-24 Jul 2021.
David Blackwell. Conditional expectation and unbiased sequential estimation. The Annals of Mathematical Statistics, 18(1):105-115 110, 1947.
Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In International Conference Learning Representations, 2016.
10

---

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. Curran Associates, Inc., 2020.
Morris L. Eaton. Multivariate statistics: : vector space approach. Wiley series in probability and mathematical statisties. Probability and mathematical statistics. Wiley, New York, 1983, ISBN 0471027766.
Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Paul Glasserman. Monte Carlo methods in financial engineering, volume 53. Springer, 2003.
Peter W. Glynn. Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, 33(10):75-88) 84, October 1990. ISSN 0001-0782.
Evan Greensmith, Peter Bartlett, and Jonathan Baxter. Variance reduction techniques for gradient estimates in reinforcement learning. In Advances in Neural Information Processing Systems, volume 14. MIT Press, 2001.
Ken Hayami. Convergence of the conjugate gradient method on singular systems. arXiv preprint arXiv:1809.00793, 2018.
Chin-Wei Huang, Ricky T. Q. Chen, Christos Tsirigotis, and Aaron Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. In International Conference on Learning Representations, 2021.
Iris A. M. Huijben, Wouter Kool, Max B. Paulus, and Ruud J. van Sloun. review of the Gumbel- max trick and its extensions for discrete stochasticity in machine learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 15(2):135311371, 2023.
M. James. The generalised inverse. The Mathematical Gazette, 62(420):109-114, 1978. ISSN 00255572, 20566328.
Martin Jankowiak and Fritz Obermeyer. Pathwise derivatives beyond the reparameterization trick. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2235-2244. PMLR, 10-15 Jul 2018.
Erik F Kaasschieter. Preconditioned conjugate gradients for solving singular systems. Journal of Computational and Applied mathematics, 24 24(152))255-275 1988.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations, 2014.
Durk Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterizzaioon trick. In Advances in Neural Information Processing Systems, volume 28, Curran Associates, Inc., 2015.
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.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, University of Toronto, Toronto, Ontario, 2009.
Brenden M. Lake, Ruslan Salakhutdinov, and Joshua B. Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266): 1 1332-1338, 2015.
Yann LeCun, Corinna Cortes, and CJ Burges. MNIST handwritten digit database. ATT Labs, 2, 2010.
Runjing Liu, Jeffrey Regier, Nilesh Tripuraneni, Michael Jordan, and Jon Mcauliffe. Rao- Blackwellized stochastic gradients for discrete distributions. In Proceedings of the 36th In- ternational Conference on Machine Learning, volume of Proceedings of Machine Learning Research, pages 4023-4031. PMLR, 09-15 Jun 2019.

---

Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte Carlo gradient estimation in machine learning. Journal of Machine Learning Research, 21(132):1-62 2020.
Jorge Nocedal and Stephen J. Wright. Numerical Optimization. Springer, 1999. ISBN 978-0-387- 98793-4.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito,, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Max B. Paulus, Chris J. Maddison, and Andreas Krause. Rao- Blackwellinnngg the straight- counbel-offtaxx gradient estimator. In International Conference Learning Representations, 2021.
M. Planitz. Inconsistent systems of linear equations. The Mathematical Gazette, 63(425):181-185 1979. ISSN 00255572, 20566328.
Robert Price. A useful theorem for nonlinear devices having Gaussian inputs. IRE Transactions on Information Theory, 4(2):69 72, 1958.
Tom Rainforth, Adam Kosiorek, Tuan Anh Lc, Chris Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 4285. PMLR, 10-15 Jul 2018.
C Radhakrishna Rao et al. Information and the accuracy attainable in the estimation of statistical parameters. Breakthroughs in statistics, pages 235 247, 1992.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the 31st International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages 1286, Bejing, China, 22-24 Jun 2014. PMLR.
Geoffrey Roeder, Yuhuai Wu, and David K Duvenaud. Sticking the landing: Simple, Lower-varianee gradient estimators for variational inference. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Ruslan Salakhutdinov and lain Murray. On the quantitative analysis of deep belief networks. In Ma- chine Learning, Proceedings of the Twenty-Fii International Conference (ICML 2008), Helsinki, Finland, June 5-9, 2008, volume 307 of ACM International Conference Proceeding Series, pages 872-879, ACM, 2008.
Tim Salimans. Implementing and automating fixed-form variational posterior approximation through stochastic linear regression. arXiv preprint arXiv: 2014.
Tim Salimans and David A Knowles. Fixed-form variational posterior approximation through stochastic linear regression. Bayesian Analysis, 8(4):837-822, 2013.
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, 2014.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In International Conference on Learning Representations, 2015.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15956):1929-1988 2014.
12

---

Michalis Titsias and Miguel Lázaro-Grre Gredilla. Doubly stochastic variational Bayes for non-conjugate inference. In Proceedings of the 31st International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages 1971-1979, Bejing, China, 22-24 Jun 2014. PMIRR
Marcin Tomczak, Siddharth Swaroop, and Richard Turner. Efficient low rank gaussian variational inference for neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, A Advances in Neural Information Processing Systems, volume 33, pages 4610-4622. Curran
Associates, Inc., 2020.
George Tucker, Dieterich Lawson, Shixiang Gu, and Chris Maddison. Doubly reparameter- ized gradient estimators for Monte Carlo objectives. In International Conference on Learning Representations, 2019.
Ronald Williams. Simple statistical gradient- following algorithms for connectionist reinforcement learning. Machine Learning. 8:229-256, 1992.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: novel image dataset for bench- marking machine learning algorithms. CoRR, abs/1708.00747. 2017.
Ming Xu, Matias Quiroz, Robert Kohn, and Scott A. Sisson. Variance reduction properties of the reparameterizzaioon trick. In Proceedings of the Tewent--econn International Conference Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 2711-2720. PMLR, 16-18 Apr 2019.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our main theoretical and experimental contributions clearly stated in the abstract and demonstrated the paper. They reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of work in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

Justification:
All proofs are provided in the technical appendices and appropriately refer-
enced.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide full experiment details in Appendix and the code to run our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
15

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes)
Justification:
We provide the code needed to run our experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ "pipsscc/puullc//uuu /otee//guudddeeeeoSuuubbmiibii ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide full experiment details in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
We report error bars in all tabulated results from our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
16

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide full details of compute resources in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with NeurlPS Code of Ethics httpp/////eaaaag..cccpppp//////a
Answer: [Yes]
Justification; After careful review of the NeurIPS Code of Ethics, it is clear that the research conducted in this paper conforms with the Code of Ethics in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper is mostly methodological. While it is possible for advances in machine learning to bring societal impacts, our paper remains general and not specific to any application so it is unlikely to bring about any immediate societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All datasets used in experiments are open-source and have been properly referenced in the paper and Appendix
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: This paper does not involve usage of LLMs.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: /ttpps////eurpps.... http:///arrrrscc///oofeeeerrr/// cc//oonferences/2225//LL. for what should or should be described.
20

---

A Supporting results
This section lists the results supporting the design of the R2-G2 estimator, namely the analytical form of conditional Gaussian distributions and the conjugate gradient algorithm.
Lemma A.1 (Eaton 1983, Proposition 3.13)) Suppose we have random vector e NNO,, I.)
where E I and linear map defined
: - T
A(x) = Ax
where A E BRxxx, Then we have the joint distribution
Onn (
I
AA
ATT]
A(e)) ~ N
mm
A
and the conditional distribution
( (A (AAT) I. - A (AAT)'A) A)
where (AAT) is the Moore- Penrose pseudo-inverse of AA
Algorithm 2 Conjugate Gradient Algorithm
Input: number of iterations T, matrix A Rmxx, column vector E I
Initialise Xg E "R
Set ro AAA Xg z.
Set Po -o--
Set 0.
while & < T and FІ2 * 0 do
r Compute ag P1 AA pt
Compute X Xx + a&Pk-
Compute r++1 = -r + AA T
PE
Compute BK+1 = FE+1
r.
Compute Pk+1 -r++1 8k+1Pk.
Increment & by 1. end while
Output: Xy.
21

---

B Proof for Proposition 4.2
We first prove that the R2-G2 estimator is an unbiased estimator of Volle [0] Recall that
R2-G2
Eio
V
R2-G2
We note that Vol
conditional expectation with W g
as the conditioning variable.
Then we have
G2
17
Rgg
Vol
Euga
Enjo
o
= Ego
Vo
Vo
[E]
where the second equality follows from the law of iterated expectations. Since the reparameterisation gradient estimator
Vol RT
an unbiased estimator of Voll.go [C]. it follows that the R2-G2 estimator
Vol R2-G2 is also an unbiased estimator of Vole.g [e].
We are left to show that the R2-G2 estimator has less variance than the reparameterisation gradient estimator. We can write
G22.22
V FolEgg
|e|| = Euga
Eijo
Vol
V Jozo
[+]
RT
27
Egga
Eigo
Vol
Euzo
<
Egg
Eijo
Vol
Vol
[01
Egg
Vil
E
where the inequality results from using Jensen's inequality, and the last equality comes from the law of iterated expectations.
C Proof of Theorem
Suppose we have a BNN linear layer where weights are independent Gaussian random variables. Given an input x E R", the pre-activations and parameters of these linear layers respectively given by
x (((("),0(())
==
E R
x g @(m)))
where (+)) = (1((0, (())),(())))))))))))))))))))(()) -() for = 1...., m. We proceed by considering fixed i.
By setting vii
(Σ)) ) 1 (diag(r()))) and W = x within Equation 6 the reparameterisation
gradient estimator given by
aē I, 12)0) o (II(E())T)
0
0z,
x
We proceed with the idea of Rao- Blackwellisation from the R2-G2 estimator by conditioning on the pre-activation z; = (1) (())) E z. This amounts to setting ( Definition 4.1
note that x (1) ( z; is equivalent to (+
x 4 (1) o., Denote 2 as
Zi
e()) (t) E with distribution denoted / Using Lemma
A.1 we can write the distribution
22

---

(00) ( as
(6) =N
( xx
In
)
Σ 2タ (o,))) )
x (0(1)))
Σラ= 2j (0()))
(1)
Define the transformation of vector to diagonal matrix as
D R Rax"
D(x) = Σ c xeLeR
where En E I is the k-th standard basis vector. Using the linearity of expectations, we have
2 x
E %
(I0))= o (,)
x²x
( 1
1n
E le E00I]
2
4
x (1D()) - 3 (1,x {))
2
2
Σ n
== 2j 01
x (x))
2
2
Σ == 23 10
Σj= Ezp (1)
(xxx)
2
Σラ= 2 (0,0)
Σ (0/))2
d 2
{xxx)
where x (x)) = (x x))
row vector with entries { and applying the reparameterisa-
tion trick gives us
2
2
2y +
Eaj
n
(of)))
& N N
Σ B
(3)
Maj 23 23
n
(o
j=l
j=l
for & ~ N (0, 1). Since we also have x"
T
Ego
In] x the Jacobian of 24 = T g( ())) is
%
Jz. (0()) = x
E
[J
go)
40
Σラ= 23 (30))) 3 {(xx0)]]
1
x
Hence, we have
R2-G2
LRT Voal
laal
This shows that the local reparameterisation gradient estimator is equivalent in distribution to Rao- Bleekklckwelleed reparameterisation gradient estimator.
23

---

We are left to show that the local reparameterisation estimator has lower variance than the global reparameterisation estimator. We have
27
LRT
R2-G2
E
Voise
Voco
[8]
Eij.s
Voiit
Voc
2
Ea.,,
Eijo
V Voiné
o( Bug
[0]
2
Eqe;
Egjo
V Voist
V (is) E
[e]
< Egg,
Enjo
Vouse
RT
V oco) E
Bu [e]
RM
1
E (c)
V
V
Oio)
18
(
where the inequality results from using Jensen's s inequality and the last equality comes from the law of iterated expectations.
D Computation of Conditional Mean
Recall that the image and kernel of a linear map T R" - RTm are defined respectively as
im(T) = (Tv v E [²) and ker(T) = {vv: Tv = 0}. Since Ae im (n(A) = im(AA ), we note
that Equation 9 has solutions in the form
(AAT)
AE + (Im (AAT) AAT y
for any y m (see James 1978);
Planitz 1979)). ). Here, (AAT)' Ac E im(AAT) im
and (In.. (AAT)* AA ly ker(AAT)
ker (T). That is, recast the matrix-vector
product (AAT)* Ae solution to Equation solution 8 to Equation it then follows that
up to additive term from ker(AT) T). For any
3" A (AAT)² AE=
24

---

Table 4: Optimisation hyperparameters and running times for experiments.
Model
Estimator
Steps
Learning Rate
Steps/s
Run Time Limit (hours)
Bayesian MLP
R2-G2
15,000
0.0001
13.57
0.5
Bayesian MLP
LRT
15,000
0.0001
48.25
0.5
Bayesian MLP
RT
15,000
0.0001
10.52
0.5
Bayesian CNN
R2-G2 12,500
0.0001
0.68
10
Bayesian CNN
RT 12,500
0.0001
2.31
5
One-layer VAE
R2-G2 100,000
0.0003
47.96
One-layer VAE
RT
100,000
0.0003
127.66
Two-layer VAE
R2-G2 100,000
0.0003
20.22
1.5
Two-layer VAE
RT 100,000
0.0003
78.72
1.5
Three-layer VAE
R2-G2
100,000
0.0003
17.11
2.5
Three-layer VAE
RT
100,000
0.0003
48.41
2.5
E
Experiment details
Compute resources
All experiments were run on single NVIDIA V100 GPU.
Architectures Model architectures for BNN and VAE experiments can be found in Section We excluded batch normalisation and dropout layers from model architectures.
Optimisation We used batch size of 80 and the Adam optimiser for all experiments Kingma and Ba (2015). We not add regularisation such as weight decay, dropout batch normalisation layers. Other optimisation parameters are listed in Table
For consistency with work on gradient estimators, report the number of optimisation/SGD steps. For BNN experiments on the MNIST and CIFAR-10 datasets, this is equivalent to training with batch size of 80 for 20 epochs.
F Comparison of Computational Cost of Gradient Estimators
In general, the worst-case complexity of computational costs from the conjugate gradient (CG) algorithm is (m,²) when one is attempting to invert a m x m dense matrix without knowing its structure. Here the worst-case would be running iterations of CG with each iteration requiring mi flops for a matrix-vector product.
In the mean-fiell setting, which is the main focus of our work, we know additional structure about the matrix A. Specifically, can factorise A WVV WVW where V Rax= is diagonal and W E Rm Moreover, we know the rank of V and the maximum rank of W, so CG only needs to be run for & min(m, n) iterations at most, by using the property of ranks that rank(AB I min(rank(AN) rarraann))))) for matrices A, B.
forward pass for linear layer using the global reparameterisation trick would use total of (2m + 1)n flops for matrix-vector products. Using the R2-G2 estimator runs at most k iterations of the CG algorithm with each iteration requiring (2m 1)n flops for matrix-vector products, so an additional (2m + 1)nk flops would be incurred at most.
For BNNs, would have m² since there complexity would be O(n) for a forward pass using the global reparamcterisation trick and an additional Omm2n) for R2-G2. For VAEs, latents to a higher-dimensions space, so the cost complexity would be ((m) for a forward pass using the global reparameterisation trick and additional (mm?) for R2-G2.
weights than pre-activations, the cost
would have m since map low-dimensional
25

---

G Additional plots of gradient variance
RT
LRT
-1.5
Log Gradient Variance
R2-G2
-2.5
-3
-3.59
-4
10
15
20
Epoch
Figure 3: Log gradient variance v.s. epoch for the bottom layer of Bayesian MLP trained on MNIST over runs. We compare the variance of gradients when training using the reparameterisation (RT), local reparameterisation (LRT) and R2-G2 estimators.
RT
R2-G2
Log Gradient Variance
-1.5
-2
-2.5
-3
10
20
30
40
Epoch
Figure 4: Log gradient variance epoch for the 8-th convolutional layer of Bayesian CNN trained on CIFAR-10 over runs. We compare the variance of gradients when training using the reparameterisation (RT) and R2-G2 estimators.
26

---

0
RT
R2-G2
Log Gradient Variance
-0.5
-1
-1.5
-2
-2.5
10
20
30
40
Epoch
Figure 5: Log gradient variance epoch for the 5-th convolutional layer of Bayesian CNN trained on CIFAR-1O over runs. We compare the variance of gradients when training using the reparameterisation (RT) and R2-G2 estimators.
H Experiments on single-layer Variational Autoencoders
Table 5: Test variational lower bounds for VAEs using the R2-G2 and Reparameterisation (RT) estimators. Higher is better. Error bars denote + 1.96 standard errors (a/ (o/v5) over 5 runs.
VAE Layers
Estimator
MNIST
Omniglot
Fashion-MNISS
R2-G2
-94.39 + 0.42
-117.61 + 2.12
-238.65 + 0.26 -238.64 64 0.25
RT
94.222 + 0.24
-117.64 2.12
The R2-G2 estimator yielded limited gains on performance for single-layer VAEs. In this setting, the R2-G2 estimator only impacts the optimisation of the encoder. This motivates experiments for hierarchical VAEs where we only apply the R2-G2 estimator within the decoder (i.e. W and V are both matrix parameters in the decoder).
I Additional plots of bounds on log-likelihood
-110
-110
-120
-120
ELBO
ELBO
130
-33
140
-140
-150
-150
20K
40K
60K
80K
100K
20K
40K
60K
80K
100K
Steps
Steps
RT: _ R2-G2
Figure 6: Bounds on log- loeliikllhood v.s. optimisation steps for two-layer VAE trained MNIST over runs. We compare the bounds on log-likelihoods when training using the reparameterisation (RT) and R2-G2 estimators on both the training set (left) and test set (right).
27

---

-130
-33
ELBO
-140
ELBO
-44
-150
-150
-160
160
20K
40K
60K
80K
100K
20K
40K
60K
80K
100K
Steps
Steps
RT:
R2-G2
Figure 7: Bounds on log-likelihood v.s. optimisation steps for a two-layer VAE trained on Omniglot over 5 runs. We compare the bounds on log-likelihoods when training using the reparameterisation (RT) and R2-G2 estimators on both the training set (left) and test set (right).
-240
-240
ELBO
ELBO
-245
-245
-250
-250
-255
-255
20K
40K
60K
80K
100K
20K
40K
60K
80K
100K
Steps
Steps
RT: - R2-G2
Figure 8: Bounds log- likeliooldo v.s. optimisation steps for two-layer VAE trained on Fashion- MNIST over runs. We compare the bounds on log-likelihoods when training using the reparameteri- sation (RT) and R2-G2 estimators both the training set (left) and test set (right).
-140
-140
ELBO
ELBO
-150
-150
-160
-60
20K
40K
60K
80K
100K
20K
40K
60K
80K
100K
Steps
Steps
RT: - R2-G2
Figure 9: Bounds on log- logelikeli v.s. optimisation steps for a three-layer VAE trained on Omniglot over runs. We compare the bounds on log-likelihoods when training using the reparametcrisation (RT) and R2-G2 estimators on both the training set (left) and test set (right).
28

---

-240
-240
ELBO
-245
ELBO
-245
-250
-250
-255
-255
20K
40K
60K
80K
100K
20K
40K
60K
80K
100K
Steps
Steps
RT: - R2-G2
Figure 10: Bounds on log-likelihood v.s. optimisation steps for three-layer VAE trained on Fashion-MNIST over 5 runs. We compare the bounds on log-likelihoods when training using the reparameterisation (RT) and R2-G2 estimators both the training set (left) and test set (right).
Licenses
Codebases:
Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization (Huang et al. 2021): MIT license.
Datasets:
MNIST (LeCun al.
2010):
Creative Commons Attribution-Sarre Alike 3.0 license
CIFAR-10 (Krizhevsky and Hinton 2009:: MIT license
Fashion-MNIst (iao et al.
2017): MIT license
Omniglot Lake et al. 2015) : MIT license
The conjugate gradient algorithm for each BNN and VAE architecture is modified from the above codebase. All experiments are performed on the above datasets.
29


Reviews:
[{"review_id": "jTws1hA04h", "reviewer": "Reviewer_zzSq", "summary": "The paper considers the following problem: estimate the gradient of an expectation that involves linear transformations of Gaussian random variables, in an unbiased and low-variance way. They propose the R2-G2 estimator, which Rao-Blackwellizes the reparameterization gradient estimator by conditioning on the linear transformation results of Gaussian variables. They prove that their estimator is unbiased and has lower variance than the reparameterization estimator. They provide closed-forms of their estimator and design an efficient (approximate) method for computing the estimator. They further prove that an existing estimator (called the local reparameterization gradient estimator) is a special case of their estimator. They empirically demonstrate that their estimator outperforms the reparameterization estimator.", "questions": "## Questions\n- I could not understand why running Algorithm 1 (and performing a backward pass) approximately computes the R2-G2 estimator as defined in Definition 4.1. More concretely, I understood the text in Section 4.2 before the paragraph \"Iterative solver\", but could not figure out why differentiating $z^*$ yields the R2-G2 estimator. Could you explain it in more detail?\n\n- Section 4.3 establishes that the local reparameterization estimator (LRT) is a special case of the R2-G2 estimator. However, I could not fully understand how much more general the R2-G2 estimator is than the LRT. Could you give concrete (preferably simple) examples where the R2-G2 is applicable but the LRT is not?\n\n## Comments\n- Section 3: The loss function $\\ell$ in Eq (1) depends on $v$ as well as $\\theta$, and this dependency looks required to express the ELBO (lines 58--59). However, in Section 3, the equations of existing estimators (SCORE, Reparameterization trick, ...) do not contain the $\\nabla_\\theta \\ell$ term. I guess we should either include the term there, or make an additional assumption that $\\ell$ does not directly depend on $\\theta$?", "limitations": "Yes, some limitations are discussed in Section 6.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- The paper tackles an important problem: develop an unbiased, low-variance estimator of the gradient of an expectation.\n- The paper proposes a novel method for the problem. Their method generalizes an existing method (local reparameterization gradient) and outperforms a popular method (global reparameterization gradient).\n- The paper is backed by both theoretical and empirical results. Theoretically, they prove the unbiasedness and a low-variance property of their method, and show a formal result that explains why an existing method (local reparameterization gradient) has performed well in practice; empirically, they demonstrate these claims in practical settings.\n- The paper is mostly well-written and well-structured.\n\n## Weaknesses\n- It was not so clear to me how their main algorithm (Algorithm 1) implements their method (Definition 4.1).\n- I did not fully understand the degree to which their method generalizes an existing method (local reparameterization gradient).", "was_revised": true, "final_justification": "I increased my score from 4 to 5 because the rebuttal addressed my main concerns.", "rebuttal": "We thank the reviewer for their review and appreciate their positive reception of our work. Please see our responses to each point raised individually below.\n\n### Weaknesses/Questions\n> It was not so clear to me how their main algorithm (Algorithm 1) implements their method (Definition 4.1).\n\n[A] We note that a forward pass consisting of reparameterisation with diagonal $\\mathbf{V} \\in \\mathbb{R}^{n \\times n}$ and linear transformation $\\mathbf{W} \\in\\mathbb{R}^{m \\times n}$ will output \n$\\mathbf{z} = \\mathbf{W} \\cdot (\\mathbf{V} \\odot \\texttt{diag}(\\epsilon))$. \n\nThis gives us the gradient\n$J_{\\mathbf{z}}(\\theta) = \\left(\\mathbf{W} \\cdot  \n    \\begin{bmatrix}\n    \\begin{array}{c|c}\n         \\mathbf{I}_n & \\frac{1}{2}\\Sigma^{-\\frac{1}{2}} \\odot (\\mathbf{1}_n \\epsilon^{\\top})\n    \\end{array}\n    \\end{bmatrix}\n    \\right)^{\\top}\n$\nin the backward pass. \n\nBy swapping $\\epsilon$ with $\\epsilon^{\\ast}$, we have a forward pass outputting $\\mathbf{z}^{\\ast} = \\mathbf{W} \\cdot (\\mathbf{V} \\odot \\texttt{diag}(\\epsilon^{\\ast}))$ with gradient\n$J_{\\mathbf{z}^{\\ast}}(\\theta) = \\left(\\mathbf{W} \\cdot  \n    \\begin{bmatrix}\n    \\begin{array}{c|c}\n         \\mathbf{I}_n & \\frac{1}{2}\\Sigma^{-\\frac{1}{2}} \\odot (\\mathbf{1}_n (\\epsilon^{\\ast})^{\\top})\n    \\end{array}\n    \\end{bmatrix}\n    \\right)^{\\top}\n$\n\nHere $J_{\\mathbf{z}^{\\ast}}(\\theta)$ is the conditional expectation of  $J_{\\mathbf{z}}(\\theta)$, using the conditional distribution $\\epsilon| \\mathbf{z}$ with mean $\\epsilon^{\\ast}$. Algorithm 1 solves for $\\epsilon^{\\ast}$ with conjugate_gradient and returns stop_gradient$(\\mathbf{z}-\\mathbf{z}^{\\ast}) + \\mathbf{z}^{\\ast}$. The return step outputs $\\mathbf{z}$ in the forward pass but uses $J_{\\mathbf{z}^{\\ast}}(\\theta)$ instead of $J_{\\mathbf{z}}(\\theta)$ in the backward pass, thereby implementing Equation 8 from Definition 4.1.\n\n\n> Section 4.3 establishes that the local reparameterization estimator (LRT) is a special case of the R2-G2 estimator. However, I could not fully understand how much more general the R2-G2 estimator is than the LRT. Could you give concrete (preferably simple) examples where the R2-G2 is applicable but the LRT is not?\n\n[B] LRT and R2-G2 estimators both use the known mean of conditional Gaussian distributions to produce Rao-Blackwellised gradients, by inverting the covariance matrix of pre-activations. As an analogy, the R2-G2 estimator generalises the LRT estimator in the same manner as inversion of a square matrix generalises inverting a scalar. We make this point more concrete below.\n\nLet $\\mathbf{v} \\in \\mathbb{R}^n$ be a vector of zero-mean independent Gaussian variables with variance parameters $(\\sigma^{(1)})^2, \\dots, (\\sigma^{(n)})^2$ contained in a diagonal covariance matrix $\\mathbf{V}$. When $\\mathbf{v}$ is mapped to a scalar output $z = \\mathbf{w}^{\\top}\\mathbf{v} \\in \\mathbb{R}$ for a vector $\\mathbf{w} \\in \\mathbb{R}^n$, the covariance matrix of $z$ is just a positive scalar $\\sigma_{z}^2 = \\sum_{i=1}^{n} w_i^2 (\\sigma^{(i)})^2$. The local reparameterisation trick exploits this convenience of working with scalars to create $m$ scalar pre-activations that are concatenated and output as a vector of Gaussian pre-activations $\\mathbf{z} \\in \\mathbb{R}^m$, by working with $m$ Gaussian vectors $\\mathbf{v}^{(1)},\\dots ,\\mathbf{v}^{(m)} \\in \\mathbb{R}^n$. The key point is that the covariance matrix of $\\mathbf{z}$ is a diagonal matrix by design to make sampling scalar pre-activations easy in the forward pass. \n\nThe subtle point here on local reparameterisation is that by square-rooting the pre-activation variance $\\sigma_{z}^2$, it actually configures reduced-variance gradients in the backward pass. The intuition here is that the square-root function $y=x^{1/2}$ has a derivative which looks like its inverse as $dy/dx = 1/2 \\times x^{-1/2}$, so applying the local reparameterisation trick in the forward pass is actually inverting the (scalar) covariance matrix $\\sigma_{z}^2$ in the backward pass. The purpose of Theorem 4.3 is to formalise this observation as Rao-Blackwellised gradients and is therefore a special case of the R2-G2 estimator. \n\nOn the point of generality of R2-G2, we note that using a diagonal covariance matrix for Gaussian pre-activations only works for BNN linear layers in the independent Gaussian weights (i.e. mean-field) setting. In general, we would use a matrix $\\mathbf{W} \\in \\mathbb{R}^{m \\times n}$ to output a pre-activation vector $\\mathbf{z} = \\mathbf{W}\\mathbf{v}$. Here $\\mathbf{z}$ has covariance matrix $\\mathbf{W}\\mathbf{V}\\mathbf{W}^{\\top}$ which is not diagonal, so LRT cannot be applied in such cases but R2-G2 can be applied. A notable example is for hierarchical VAE architectures where $\\mathbf{v}$ is a vector of latent Gaussian variables used in the decoder and $\\mathbf{W}$ is the linear transformation that immediately follows it (please see our VAE experiments in Section 5.3).\n\nTo summarise, LRT and R2-G2 estimators both invert the covariance matrix of Gaussian pre-activations to produce reduced-variance gradients. LRT inverts a (scalar) matrix automatically as a result of the parameterisation used in the forward pass, but the application is restricted to BNN linear layers. On the other hand, R2-G2 inverts a symmetric square matrix after observing pre-activations in the forward pass, by using the conjugate gradient algorithm as a matrix inversion procedure and thereby making it more general than LRT. To further motivate the R2-G2 estimator, we propose to include this explanation at the start of Section 4 in a potential camera-ready version.\n\n> Section 3: The loss function $\\ell$ in Eq (1) depends on $\\mathbf{v}$ as well as $\\theta$, and this dependency looks required to express the ELBO (lines 58--59). However, in Section 3, the equations of existing estimators (SCORE, Reparameterization trick, ...) do not contain the term $\\nabla_{\\theta} \\ell$. I guess we should either include the term there, or make an additional assumption that does not directly depend on $\\theta$?\n\n[C] Thank you for spotting the typo in the ELBO (lines 58--59). We agree that the dependency on $\\theta$ can be better clarified here. In the next version, we propose to write $\\ell_{\\mathcal{D}, f}(\\mathbf{v})$ in the Equation after Line 51 and $\\ell(\\mathbf{v})$ in the Equation after Line 63. This would better describe the problem setting without assuming $q$ is reparameterisable in Section 2. We also propose to only make an explicit dependency on $\\theta$ in Section 3 when $q$ is assumed to be reparameterisable. Please also see our answer [C] to reviewer j5Ym.\n\nTo clarify the equations in Section 3, we note that the gradient term $\\nabla_{\\theta} \\ell$ is generally unknown and is exactly why we need gradient estimators! We have followed the hat notation from statistics to note that equations in Section 3 are estimators of $\\nabla_{\\theta} \\ell$."}, {"review_id": "jIEOLJu0do", "reviewer": "Reviewer_j5Ym", "summary": "A common problem in probabilistic machine learning is to learn the parameters of a distribution q used to integrate out latent random variables from an objective function. In other words, the goal is to optimize the expectation of an objective function with respect to the distribution q by adjusting the parameters of q.  Because the expectation depends on q, the gradient cannot simply be pushed inside the expectation. For example, in Bayesian inference, the latent variables are the parameters of the model and the distribution q is the variational approximation to the posterior. In VAEs, the latent variables are the encodings and the distribution q is the encoder.\n\nThe authors start by discussing three common approximations, presented in decreasing order of variance: REINFORCE, which uses the log derivative trick, the reparameterization trick (RT), which requires q to be reparameterizable, and the local reparameterization trick (LRT), which applies when the latent variable can be expressed as an affine transformation of Gaussian random variables (for example, the pre-activations in a neural network). The local reparameterization trick tends to have the least variance, but also requires the latent variables to factorize under q, which happens, for example, with a mean-field Gaussian Bayesian neural network. \n\nThis paper introduces a new estimator, R2-G2, designed to acheive low variance like the local reparameterization trick but can apply when q does not factorize. The R2-G2 estimator performs Rao Blackwellization on the RT estimator by conditioning on a single sample of the pre-activations. As with any Rao Blackwellization, this means it will have lower variance. Note that the second in G R2-G2 is for Gaussian because their estimator applies to Gaussian q.\n\nThe paper then discusses computing their estimator by solving a least squares problem and shows that it is equivalent to the LRT for a mean-field Gaussian BNN. \n\nThere are two experiments in the paper: image classification with BNNs and hierarchical VAEs, both on standard benchmark datasets. For the BNN, the authors compare against RT and LRT on MNIST, and against RT on CIFAR-10. They do not compare against LRT on CIFAR-10 because they use convolutional layers, which the LRT does not support because q would not factorize. Looking at the results in Table 2, the confidence intervals appear to all overlap. For the VAE experiment, the authors compare against RT and find that R2-G2 has better ELBO value on test data.", "questions": "Questions:\n- Could a non-mean-field example be run for the BNN, which would allow R2-G2 to be tested in a situation where the LRT doesn’t apply? \n- Is the CNN example main architecture where R2-G2 applies but LRT does not? What about attention? \n- Is there a reason why the LRT is not run for the VAE example? It is a fully factorized Gaussian variational posterior so my understanding is that it can be applied. \n- Is computing $\\beta^*$ practical for large architectures? How long did you need to run conjugate gradients? I did not see a description of the MLP architecture in the paper (the appendix references the main text and the main text cites Srivastava 2014, but this paper has many MLP examples).  \nNote the first 3 questions are about understanding where R2-G2 is most useful. \n\nSmall:\n- MLP, ELBO (in intro) not defined \n- The equation after line 63 seems to assume q can be reparameterized\n- Notation: Is f parameterized by v? \n- The equation after line 63 appears to assume $q$ is reparameterizable", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "The paper is clear and well written. The problem addressed by the paper is a generic and common problem in ML so any theoretical or practical improvements are significant. \n\nRegarding originality, the paper is conceptually straightforward as it applies a well-known technique for variance reduction (Rao Blackwellisation) to a well-known estimator (RT). Proposition 4.2 seems like an expected result. Theorem 4.3 is more interesting to me because it provides an interpretation of LRT. \n\nTo me there are two questions the experiments should investigate: (1) is the estimator meaningfully reducing variance? (2) is the estimator computationally feasible when it is? \n\nThe latter question is not really addressed in the experiments since the architectures are small, I believe (though I am not sure what they are exactly, see question below).\n\nRegarding first question, for the MNIST dataset the LRT and R2-G2 should be the same by Theorem 4.3 (which appears to be the case) but they have little to no improvement over RT. I would be interested in seeing a full covariance variational family (or something not factorized), since in this setting (as I understand it) R2-G2 is not guaranteed to be the same as LRT. Similarly for CNN on CIFAR-10, R2-G2 and RT have similar performance, but LRT doesn’t apply here. \n\nFor the VAE, I didn’t follow why LRT was not used (see mu question below). It is nice to see that R2-G2 had a larger improvement over RT than in the BNN case, but the confidence intervals were overlapping. \n\nIn general, my primary concern is that the experiments don’t seem to verify the issues of the LRT discussed in Section 4, namely performance like the LRT, better than RT, but in situations where it can’t be used (unless it can’t be used for the VAE).", "was_revised": true, "final_justification": "Given that this method requires additional computation beyond RT, ideally there would be stronger empirical results, especially in the BNN example (as I interpret the results, all methods are about the same). However, the interpretation of LRT as Rao-Blackwellisation is a contribution on its own and there is stronger evidence that R2-G2 improves over RT in the VAE example where LRT does not apply (clarified to me during the rebuttal). \n\nFor what it is worth, the name “R2-G2” was not a factor in my review, but I did appreciate it.", "rebuttal": "We thank the reviewer for their review and appreciate their engagement with our work. Please see our responses to each question raised individually below.\n\n### Questions\n\n> Could a non-mean-field example be run for the BNN, which would allow R2-G2 to be tested in a situation where the LRT doesn’t apply?\n\n> Is the CNN example main architecture where R2-G2 applies but LRT does not? What about attention?\n\n> Is there a reason why the LRT is not run for the VAE example? It is a fully factorized Gaussian variational posterior so my understanding is that it can be applied.\n\n[A] Great questions! We note that the R2-G2 estimator is only applicable where the reparameterisation trick is used to sample Gaussian variables, we do not include attention layers since standard attention layers do not usually apply the reparameterisation trick to sample Gaussian variables. We also note that using the global reparameterisation trick on non-mean-field BNNs would introduce a separate scalability challenge that would be difficult to tackle within this work. To see this, we can take a BNN linear layer has a $m \\times n$ matrix of Gaussian weights as an example. This would require parameterising a dense $mn \\times mn$ covariance matrix (or its Cholesky factor as a lower triangular matrix), making computations challenging for both the reparameterisation trick and the R2-G2 estimator by extension.\n\nMean-field BNN convolutional layers and VAEs are both examples where LRT cannot be applied. To clarify the VAE example, the fully factorised Gaussian assumption is applied in the latent space $\\mathbf{v}$ so the reparameterisation trick applied here is the global reparameterisation trick and RT estimator. That is, the local reparameterisation trick and LRT estimator are not usable for VAEs because they work by summing over multiple independent Gaussian variables that are mapped to scalar Gaussian pre-activations. The linear transformation of these VAE latents $\\mathbf{W}\\mathbf{v}$ (i.e. pre-activations) does not have a factorised Gaussian variational posterior in general, so LRT cannot be applied. We make this more precise below.\n\nLet $\\mathbf{v} \\in \\mathbb{R}^n$ be a vector of zero-mean independent Gaussian variables with variance parameters $(\\sigma^{(1)})^2, \\dots, (\\sigma^{(n)})^2$ contained in a diagonal covariance matrix $\\mathbf{V}$. When $\\mathbf{v}$ is mapped to a scalar output $z = \\mathbf{w}^{\\top}\\mathbf{v} \\in \\mathbb{R}$ for a vector $\\mathbf{w} \\in \\mathbb{R}^n$, the covariance matrix of $z$ is just a positive scalar $\\sigma_{z}^2 = \\sum_{i=1}^{n} w_i^2 (\\sigma^{(i)})^2$. The local reparameterisation trick exploits this convenience of working with scalars to create $m$ scalar pre-activations that are concatenated and output as a vector of Gaussian pre-activations $\\mathbf{z} \\in \\mathbb{R}^m$, by working with $m$ Gaussian vectors $\\mathbf{v}^{(1)},\\dots ,\\mathbf{v}^{(m)} \\in \\mathbb{R}^n$. The key point is that the covariance matrix of $\\mathbf{z}$ is a diagonal matrix by design to make sampling scalar pre-activations easy in the forward pass. \n\nThe subtle point here on local reparameterisation is that by square-rooting the pre-activation variance $\\sigma_{z}^2$, it actually configures reduced-variance gradients in the backward pass. The intuition here is that the square-root function $y=x^{1/2}$ has a derivative which looks like its inverse as $dy/dx = 1/2 \\times x^{-1/2}$, so applying the local reparameterisation trick in the forward pass is actually inverting the (scalar) covariance matrix $\\sigma_{z}^2$ in the backward pass. The purpose of Theorem 4.3 is to formalise this observation as Rao-Blackwellised gradients and is therefore a special case of the R2-G2 estimator. \n\nWe note that using a diagonal covariance matrix for Gaussian pre-activations only works for BNN linear layers in the independent Gaussian weights (i.e. mean-field) setting. In general, we would use a matrix $\\mathbf{W} \\in \\mathbb{R}^{m \\times n}$ to output a pre-activation vector $\\mathbf{z} = \\mathbf{W}\\mathbf{v}$. Here $\\mathbf{z}$ has covariance matrix $\\mathbf{W}\\mathbf{V}\\mathbf{W}^{\\top}$ which is not diagonal, so LRT cannot be applied in such cases but R2-G2 can be applied. A notable example is for hierarchical VAE architectures where $\\mathbf{v}$ is a vector of latent Gaussian variables used in the decoder and $\\mathbf{W}$ is the linear transformation that immediately follows it (please see our VAE experiments in Section 5.3).\n\nTo summarise, LRT and R2-G2 estimators both invert the covariance matrix of Gaussian pre-activations to produce reduced-variance gradients. LRT inverts a (scalar) matrix automatically as a result of the parameterisation used in the forward pass, but the application is restricted to BNN linear layers. On the other hand, R2-G2 inverts a symmetric square matrix after observing pre-activations in the forward pass, by using the conjugate gradient algorithm as a matrix inversion procedure and thereby making it more general than LRT. To further motivate the R2-G2 estimator, we propose to include this explanation at the start of Section 4 in a potential camera-ready version. We agree that an even more general non-mean-field setting (where $\\mathbf{V}$ is not diagonal) will be a very interesting avenue for future work and further experiments!\n\n> Is computing $\\beta^{\\ast}$ practical for large architectures? How long did you need to run conjugate gradients? I did not see a description of the MLP architecture in the paper (the appendix references the main text and the main text cites Srivastava 2014, but this paper has many MLP examples).\n\n[B] The computation of $\\beta^{\\ast}$ is done for each layer where the R2-G2 estimator is applied. This would be practical in the top/final layers of neural network architectures, even for larger architectures, for two reasons. First, these layers are where features are compressed to low-dimensions, usually with a linear layer, so this should limit the number of iterations of conjugate gradients. Second, the variance of gradients of top/final layers are usually much higher than bottom/hidden layers (we can see this in the plots of RT gradient variances in Figures 1 and 3), so there is more benefit to Rao-Blackwellising gradients in top/final layers when training. \n\nIn our experiments, we ran conjugate gradient (CG) using the property of ranks that $\\text{rank}(AB) \\leq min(\\text{rank}(A), \\text{rank}(B))$ for matrices $A,B$ of appropriate shapes. For each NN layer where the R2-G2 is used, this equates to running maximum of: 1 CG iteration for mean-field BNN linear layers, 4 to 16 CG iterations for mean-field BNN convolutional layers, 50 CG iterations for VAE sampling layers in the decoder.\n\nThe MLP architecture we used from Srivastava 2014 was a MLP with two hidden layers of 1024 ReLU units. We propose to explicitly state this detail in Section 5.2 of the Experiments section in a potential camera-ready version.\n\nSmall questions\n> MLP, ELBO (in intro) not defined\n> The equation after line 63 seems to assume q can be reparameterized\n> Notation: Is f parameterized by v?\n> The equation after line 63 appears to assume is reparameterizable\n\n[C] Thanks for these comments as it will improve clarity of our writing. In the next version, we propose to introduce MLP within the Experiments section and ELBO in Line 57, and to write $\\ell_{\\mathcal{D}, f}(\\mathbf{v})$ in the Equation after Line 51 and $\\ell(\\mathbf{v})$ in the Equation after Line 63 to describe the problem setting without assuming $q$ is reparameterisable in Section 2."}, {"review_id": "63aTEVkBYC", "reviewer": "Reviewer_jPqY", "summary": "This paper proposes the R2-G2 gradient estimator. \nThis gradient estimator generalizes the local reparameterization gradient by highlighting connections to Rao-Blackwellization and proposes an implementation for the general case of models with latent Gaussian variables. \nThe estimator is evaluated against the reparameterization gradient estimator in hierarchical variational auto-encoders and bayesian neural networks.", "questions": "- Can you expand discussion on the computational costs of the conjugate gradient method relative to the full forward and backward pass of the model? \n- Can the authors clarify the backward pass? Does their implementation differentiate through the roll-outs conjugate gradient algorithm or otherwise?\n- Can you provide results comparing R2-G2 and RT on the same compute budget to provide an equal playing field?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Quality:\n- The paper highlights some interesting connections between the local reparam trick and Rao-Blackwellization and R2-G2 naturally emerges as a generalization of the local reparam gradient as a result.\n- The technical exposition in the paper is very clear, all claims are well-supported and results are derived. \n- Kudos to the authors for being transparent and honest about the experimental results. \n- Unfortunately, the method seems rather compute-intensive (running conjugage-gradient algorithm as part of the forward-pass), which makes me question how practically relevant it is. \n\nClarity\n- This paper is extremely well-written and a pleasure to read. \n- The experimental exposition given is thorough. \n- The mathematical notation used is consistent and clear. \n\nSignificance\n- I do believe that the theoretical analysis and insights into local reparam gradients, connections the paper draws to Rao-Blackwellization, and the generalized R2-G2 estimator the paper proposes are valuable to the research community.  \n- My main criticism is with the benefit of using the estimator in practise. The improvements in some cases are marginal (e.g., BNN on MNIST, CIFAR-10, no improvement over RT) or non-existent (VAE, one layer, as the authors acknowledge).  However, I appreciate the transparency and honesty about the experimental results. BNN on MNIST, CIFAR-10 which is equivalent to the method, also got published and doesn't do better in these settings than RT (as it's equivalent R2-G2). The variance reduction is clearly visible, but doesn't translate into improved performance, possibly because the tasks are too easy.  Further, the experimental evaluation is limited to \"initial training\". However, the authors make a good argument for motivating this setting, which I appreciate.\n\nOriginality\n- In my opinion, the R2-G2 estimator provides a valuable new perspective on local reparm gradients. Rao-Blackwellization is of interest in the gradient estimation community too, so this work is a natural addition to existing work.", "was_revised": true, "final_justification": "I read the rebuttal and my questions were clarified. The concern about the practical significance and the empirical improvements of this method remain. However, I still believe that the reasons to accept outweigh the reasons to reject, so I am keeping my score as is.", "rebuttal": "We thank the reviewer for their review and appreciate their positive reception of our work. Please see our responses to each question raised individually below.\n\n### Questions\n\n> Can you expand discussion on the computational costs of the conjugate gradient method relative to the full forward and backward pass of the model?\n\n[A] Great question! We first note that the worst-case complexity of computational costs from the conjugate gradient (CG) algorithm is $\\mathcal{O}(m^3)$ (line 188) when one is attempting to invert a $m \\times m$ dense matrix $C$ without knowing its structure. Here the worst-case would be running $m$ iterations of CG with each iteration requiring $m^2$ flops for a matrix-vector product.\n\nIn the mean-field setting, we know more structure about the matrix $C$ and can factorise it as $C = \\mathbf{W}\\mathbf{V}\\mathbf{W}^{\\top}$ where $\\mathbf{V} \\in \\mathbb{R}^{n \\times n}$ is diagonal and $\\mathbf{W} \\in \\mathbb{R}^{m \\times n}$. Moreover, we know the rank of $V$ and the maximum rank of $\\mathbf{W}$, so CG only needs to be run for $k=min(m,n)$ iterations at most, by using the property of ranks that $\\text{rank}(AB) \\leq min(\\text{rank}(A), \\text{rank}(B))$ for matrices $A,B$ of appropriate shapes.\n\nA forward pass for a linear layer using the global reparameterisation trick would use a total of $(2m + 1)n$ flops for matrix-vector products. Using the R2-G2 estimator runs at most $k$ iterations of the CG algorithm with each iteration requiring $(2m + 1)n$ flops for matrix-vector products, so an additional $(2m + 1)nk$ flops at most. \n\nIn our paper, we have used $n$ and $m$ to refer to the dimensions of the input and output/pre-activations Gaussian variables. For BNNs, we would have $n > m^2$ since there are more weights than pre-activations, so the cost complexity would be $\\mathcal{O}(n)$ for a forward pass using the global reparameterisation trick and an additional $\\mathcal{O}(n)$ for R2-G2. For VAEs, we would have $m > n$ since we map low-dimensional latents to a higher-dimensional space, so the cost complexity would be $\\mathcal{O}(m)$ for a forward pass using the global reparameterisation trick and an additional $\\mathcal{O}(mn)$ for R2-G2. In the next version, we propose to include this discussion in the Appendix.\n\n> Can the authors clarify the backward pass? Does their implementation differentiate through the roll-outs conjugate gradient algorithm or otherwise?\n\n[B] Our implementation, presented as Algorithm 1, does not differentiate through $\\beta^{\\ast}$ or any of the roll-outs/iterates generated by the conjugate gradient algorithm. Once $\\beta^{\\ast}$  and $\\epsilon^{\\ast}$ are computed, they are only treated as constant vectors. Specifically, $\\epsilon^{\\ast}$ is used to compute $\\mathbf{z}^{\\ast} = \\mathbf{W} \\cdot (\\mathbf{V} \\odot \\texttt{diag}(\\epsilon^{\\ast}))$ so that the Jacobian from Equation 8 within Definition 4.2 is used for backpropagation. Please also see our answer [A] to reviewer zzSq for a detailed explanation.\n\n> Can you provide results comparing R2-G2 and RT on the same compute budget to provide an equal playing field?\n\n[C] While a comparison can theoretically be made to compare R2-G2 and RT based on the number of FLOPs used for both estimators, we note that this is difficult to track in practice since we do not know the following beforehand: convergence behaviour of the conjugate gradient algorithm on the problem at hand and the exact FLOPs used before it terminates.\n\nIn our hierarchical VAE comparisons, we fixed a small number of gradient steps (100K) and a limit in run-time on the same computing resource (a single NVIDIA V100 GPU). The run-time limit was set to 90 minutes for two-layer VAEs and 150 minutes for three-layer VAEs. Empirically, we observed that R2-G2 experiment runs were all completed in these limits with no more than 20 minutes of additional run-time compared to RT experiment runs. We propose to include these run-time limits and empirical observation within the Experiments section in a potential camera-ready version."}, {"review_id": "Gai95zFUQC", "reviewer": "Reviewer_ZPWE", "summary": "The authors discuss the problem of estimating stochastic gradients of a function, specifically `\\nabla_\\theta \\mathbb{E}_{q_\\theta}[f(\\theta, \\dots)]`, which are omnipresent in many applications of approximate inference for neural networks. For this, the authors revise the reparameterisation trick and present a Rao-Blackwellised reparameterisation applicable if $q_\\theta$ (e.g., the variational family) follows a mean-field Gaussian distribution. This new reparameterisation is then shown to be equivalent to *local* reparameterisation under certain models and empirically (as well as theoretically) reduces variance in the gradient estimates compared to standard reparameterisation. The authors also show some slight improvements in performance on some basic image classification tasks (MNIST and CIFAR-10).", "questions": "- Why make the assumption that $v$ is mean-field Gaussian, if you could also derive the distribution for correlated $v$?\n- What is the effect of making all of those independence assumptions?\n- What is $\\text{im}(X)$ and $\\ker(X)$?\n- What is the computational cost in FLOPS?\n- How does this approach scale and perform in more interesting settings (e.g., ImageNet)?\n- Why doesn't the convergence improve?", "limitations": "The authors discuss some of the limitations of the work in the conclusion section.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\nThe idea is novel and methodologically sound. The work involves various rather technical steps, and I appreciate the effort in trying to make the approach scalable. The approach clearly enjoys variance reduction and is convincing in this regard. \n\n### Weaknesses\nUnfortunately, the paper is not very well written. Many parts feel unclear at first, and the reader needs to fill in the gaps to understand what the authors are doing or trying to convey. This makes the reading rather unpleasant.\n\nSecondly, the performance improvements are marginal at best. I don't think this is a problem, as I would not necessarily expect big performance gains, but I was surprised to see that the convergence is not improved even though the variance is substantially reduced. This makes me wonder to what extent the high additional complexity is justified, as little gains can effectively be observed.\n\nLastly, the evaluation is only on toy data by today's standards. I would have expected to see an evaluation in more interesting contexts.\n\n### Additional Comments\n\nBelow is a list of additional comments that are not necessarily questions.\n\n- Line 14, the authors say that \"latent variables [...] enable us to embed inductive biases like uncertainty in the model specification\". I do not understand what this is meant to say. I agree that LVs can help encode prior assumptions, but those aren't necessarily linked to the model's inductive biases. I also do not understand how uncertainty in the model specification is an inductive bias. I see that the inductive bias will affect the distribution of uncertainty over the model specification, but this paragraph seems somewhat convoluted.\n- Line 16, the authors say LV appears in \"modelling tasks, including Bayesian inference\". I do not see how Bayesian inference is a modelling task.\n- I generally found the distinction between estimators and \"single-sample\" estimators confusing. While the reparameterization trick is often used with a single Monte-Carlo sample, I do not see that there is a need to distinguish it from other estimators, as this is a hyperparameter. It is very much conceivable to use more than one sample.\n- Eq. 3 inherently assumes that $v \\sim g(\\epsilon^{(i)}, \\theta^{(i)})$ with $v \\in \\mathbb{R}^n$ is distributed by a mean-field Gaussian. This is a rather strong assumption that is not made explicit. As for the current version, this means the authors assume independence in $v$ and in $z$, which seems very strong.\n- Line 102, $\\tau$ is not introduced, and it is unclear how this relates to Eq. 3. Additional confusion happens later on where $\\tau$ is used.\n- Line 121, typo \"is\" -> \"are\"\n- Paragraph 118 - 124, I believe it would be worth mentioning that the \"global\" reparameterisation trick is in fact not limited to the Gaussian family but can be applied through inverse CDF sampling to any distribution.\n- Section 4.2., I had a really hard time piecing things together in this section. For example, $\\text{im}(X)$ and $\\ker(X)$ are never introduced. The solution of Eq.9 trivially results from Planitz, but it is hard to understand how this now relates to the goal we are after here. Moreover, it is unclear if $\\epsilon^\\star$ refers to the mean of the distribution over $\\epsilon$ as in Sec. 4.1 or not, as it is referred to as the \"fitted noise\" and $\\beta^\\star$ is not the mean of the distribution of $\\beta$ but the solution to the quadratic program. \n- I believe a more detailed cost estimate of the R2-G2 estimator is necessary to fully characterise the computational cost (ideally in O notation and FLOPS) and compare to vanilla reparameterisation.", "was_revised": true, "final_justification": "The authors have sufficiently addressed my questions and concerns.", "rebuttal": "We thank the reviewer for their review and appreciate their constructive feedback on our work. Please see our responses to each point raised individually below.\n\n### Weaknesses\n> Many parts feel unclear at first.\n\nPlease see our proposed changes in the responses below. We hope these changes improve the reading of our work and make it more accessible to readers.\n\n### Comments\n> Line 14, the authors say that \"latent variables [...] enable us to embed inductive biases like uncertainty in the model specification\". I do not understand what this is meant to say. \n> Line 16, the authors say LV appears in \"modelling tasks, including Bayesian inference\". I do not see how Bayesian inference is a modelling task.\n\n[A] We agree that the language can be expressed better here. In the next version, we propose to change line 14 to \"latent variables [...] enable us to embed prior assumptions into models, such as uncertainty in model parameters [...]\" and line 16 to “modelling tasks, including variational inference  [...]”.\n\n> I generally found the distinction between estimators and \"single-sample\" estimators confusing. While the reparameterization trick is often used with a single Monte-Carlo sample, I do not see that there is a need to distinguish it from other estimators, as this is a hyperparameter. It is very much conceivable to use more than one sample.\n\n[B] We note the distinction between single-sample estimators and more general (potentially multi-sample) estimators is crucial on two fronts.\n\nFirst, the distinction helps distinguish work on single-sample gradient estimators, such as [1] and our work, from the wide literature on gradient estimators that use multiple samples and control variates to compute the expected gradient.\n\nSecond, it clarifies where Rao-Blackwellisation is actually applied. One application of R2-G2 will only reduce the gradient variance from one set of Gaussian variables. While applicable to each sample in a mini-batch, it remains a sample-level procedure that is parallelised. On the other hand, multi-sample estimators focus on the expected gradient, so Rao-Blackwellising this implies conditioning over a sample of $K>1$ pre-activation vectors per sample, which is not what we do.\n\n> Line 102, $\\tau^{(i)}$ is not introduced, and it is unclear how this relates to Eq. 3. Additional confusion happens later on where $\\tau^{(i)}$ is used.\n\n[C] $\\tau^{(i)}$ defines a vector of variance parameters in the same way as $\\tau$ introduced in Line 86. For clarity, we propose to explicitly state $\\tau^{(i)} = ((\\sigma_{1}^{(i)})^2, \\dots, (\\sigma_{n}^{(i)})^2)$ within the local reparameterisation trick paragraph in Section 3.\n\n> Paragraph 118 - 124, I believe it would be worth mentioning that the \"global\" reparameterisation trick is in fact not limited to the Gaussian family but can be applied through inverse CDF sampling to any distribution.\n\nGreat suggestion! We propose to add this in the next version.\n\n### Questions\n> Why make the assumption that is mean-field Gaussian, if you could also derive the distribution for correlated?\n> What is the effect of making all of those independence assumptions?\n\n[D] Great question! We have focused on the independent case in this work for two reasons. \nFirst, we wanted to establish a first connection between the classical idea of Rao-Blackwellisation from statistics and the local reparameterisation trick (which enjoys low-variance gradients in practice but can only be derived for mean-field BNN linear layers). Second, we wanted to assess the benefits and limitations from training with Rao-Blackwellised gradients would benefit probabilistic models in this simple setting before investigating the more difficult setting of correlated Gaussians in future work. This uncovered sizeable gains for hierarchical VAEs and limitations for single-layer VAEs when training them with R2-G2. We agree that the independence assumption is not necessary to derive the R2-G2 estimator and investigating the correlated noise setting serves as an interesting direction for future work! We note that the independence assumption affects model specification, which affects training with the RT estimator as much as it does the R2-G2 estimator.\n\nIn our work, we have presented the R2-G2 estimator under the assumption of independent Gaussians to present the clear connection to Rao-Blackwellisation as it is one of our key contributions. While we agree that independent assumption is strong, we note that it is also a common assumption when evaluating a test bed of standard probabilistic models (e.g. VAEs, BNNs). We propose to state that the expression for the R2-G2 estimator presented in Definition 4.1 only holds under the independent Gaussian assumption and that a more general expression can be derived for the correlated Gaussian setting in a potential camera-ready version.\n\n\n>Moreover, it is unclear if $\\epsilon^{\\ast}$ refers to the mean of the distribution over $\\epsilon$ as in Sec. 4.1 or not, as it is referred to as the \"fitted noise\" [...].\n> What is im(X) and ker(X)?\n\n[E] $\\texttt{im}(\\mathbf{X})$ and $\\texttt{ker}(\\mathbf{X})$ refer to the image and kernel of the matrix $\\mathbf{X}$ respectively. We agree that introducing these quantities in Section 4.2 affects the readability of the section. In the next version, we propose to defer these quantities to the Appendix to improve readability of the paper.\n\nIndeed $\\epsilon^{\\ast}$ refers to both the mean of the conditional distribution $\\epsilon| \\mathbf{W} \\cdot g = \\mathbf{z}$ from Section 4.1 and the fitted noise from Line 178. Using the language of linear models, fitting a linear model (Equation after line 172) is the same as mapping “observed values” to “fitted values”. For consistency with this language, we propose to change Line 178 to “projecting the observed noise $\\epsilon$ to the fitted noise $\\epsilon^{\\ast}$ [...]”.\n\n> What is the computational cost in FLOPS?\n> How does this approach scale and perform in more interesting settings (e.g., ImageNet)?\n\n[F] Great question! We first note that the worst-case complexity of computational costs from the conjugate gradient (CG) algorithm is $\\mathcal{O}(m^3)$ (line 188) when one is attempting to invert a $m \\times m$ dense matrix $C$ without knowing its structure. Here the worst-case would be running $m$ iterations of CG with each iteration requiring $m^2$ flops for a matrix-vector product.\n\nIn the mean-field setting, we know more structure about the matrix $C$ and can factorise it as $C = \\mathbf{W}\\mathbf{V}\\mathbf{W}^{\\top}$ where $\\mathbf{V} \\in \\mathbb{R}^{n \\times n}$ is diagonal and $\\mathbf{W} \\in \\mathbb{R}^{m \\times n}$. Moreover, we know the rank of $V$ and the maximum rank of $\\mathbf{W}$, so CG only needs to be run for $k=min(m,n)$ iterations at most, by using the property of ranks that $\\text{rank}(AB) \\leq min(\\text{rank}(A), \\text{rank}(B))$ for matrices $A,B$.\n\nA forward pass for a linear layer using the global reparameterisation trick would use a total of $(2m + 1)n$ flops for matrix-vector products. Using the R2-G2 estimator runs at most $k$ iterations of the CG algorithm with each iteration requiring $(2m + 1)n$ flops for matrix-vector products, so an additional $(2m + 1)nk$ flops at most. \n \nIn our paper, we have used $n$ and $m$ to refer to the dimensions of the input and output/pre-activations Gaussian variables. For BNNs, we would have $n > m^2$ since there are more weights than pre-activations, so the cost complexity would be $\\mathcal{O}(n)$ for a forward pass using the global reparameterisation trick and an additional $\\mathcal{O}(n)$ for R2-G2. For VAEs, we would have $m > n$ since we map low-dimensional latents to a higher-dimensional space, so the cost complexity would be $\\mathcal{O}(m)$ for a forward pass using the global reparameterisation trick and an additional $\\mathcal{O}(mn)$ for R2-G2. In the next version, we propose to include this discussion in the Appendix.\n\nWe note that using R2-G2 would be practical in the top/final layers of neural network architectures, even for larger architectures, for two reasons. First, these layers are where features are compressed to low-dimensions, usually with a linear layer, so this limits the number of iterations of conjugate gradient. Second, the variance of gradients of top/final layers are much higher than bottom/hidden layers (we can see this in the plots of RT gradient variances in Figures 1 and 3), so there is more benefit to Rao-Blackwellising gradients in top/final layers when training.\n\n> Why doesn't the convergence improve?\n\n[G] For BNNs, training with reduced-variance gradients (R2-G2) does not exhibit large gains over RT. We believe that gains in BNNs are more driven by a good architecture choice to the dataset, such as using a CNN for CIFAR-10. Please see the comments by Reviewer jPqY and Figure 1a of the original local reparameterisation trick paper [2] which shows similar error rates between the local and global reparameterisation tricks (referred to as Variational (A) and Variational (B) respectively in [2]).  We also note that the architectures have been simplified in our experiments, as we do not use dropout or batch normalisation layers. While this can have an effect on model performance, we wanted to better observe the variance reduction from R2-G2 in isolation (see Figures 1, 3, 4, 5) and whether it yields performance gains. \n\nFor hierarchical VAEs, we show in Table 3 and Figure 2 that training the decoder improves the convergence of hierarchical VAE models. In the case of three-layer VAEs on MNIST and Omniglot datasets, we see an improvement of over 9 nats and 1 nat respectively which is substantial.\n\n* [1] Max B Paulus, Chris J. Maddison, and Andreas Krause. “Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator”. In ICLR 2021.\n* [2] Durk P Kingma, Tim Salimans, and Max Welling. “Variational dropout and the local reparameterization trick”. In NeurIPS 2015."}]

Author rebuttal:

Meta-review: The paper proposes a reparameterization gradient for Monte Carlo gradient estimation with lower variance compared to existing approaches. The paper is clearly written, the approach is novel, and the method is supported by strong numerical experiments. That said, there are some limitations to the method such as higher computational cost, and the contribution is on the more incremental side.
Decision: Accept (poster)