Title: Rao-Blackwellised Reparameterisation Gradients
Venue: NeurIPS 2025
Accepted: yes
Keywords: reparameterisation trick

stochastic gradient estimation

Rao-Blackwellisation

variance reduction
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors discuss some of the limitations of the work in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nQuality:\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.\n\nQUESTIONS:\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- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nA 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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).\n\nQUESTIONS:\nQuestions:\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\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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).\n\nQUESTIONS:\n## 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$?\n\nLIMITATIONS:\nYes, some limitations are discussed in Section 6.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}