Research map / Probabilistic and scientific modeling / Variational and generative sampling
Rao-Blackwellised Reparameterisation Gradients
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes R2-G2, a single-sample gradient estimator that Rao-Blackwellises the reparameterisation gradient by conditioning on linear pre-activations of Gaussian variables (solved via conjugate gradient), showing the local reparameterisation trick is a special case.
Problem
High variance of reparameterisation gradients for latent Gaussian variables in models where the local reparameterisation trick is not applicable, such as when pre-activations do not have a diagonal covariance.
Approach
Condition the reparameterisation noise on linear transformations (pre-activations) of Gaussian variables, yielding a Rao-Blackwellised single-sample gradient estimator computed by solving a least-squares/conjugate-gradient problem; the local reparameterisation trick is shown to be a special case for independent-weight linear layers.
Key results
- On three-layer hierarchical VAEs, R2-G2 improves test bounds over RT by 9.05 nats on MNIST and 1.17 nats on Omniglot.
- On Bayesian MLP for MNIST, R2-G2 and LRT give matching log-likelihoods and gradient variances, both lower than RT.
- On Bayesian CNN for CIFAR-10, R2-G2 gives higher test log-likelihood than RT (-3.87 vs -3.88) with similar accuracy (67.97% vs 67.98%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace RT with R2-G2 | Bayesian MLP · MNIST | test log-likelihood: -3.07 ± 0.00 -> -3.06 ± 0.00 | Helped |
| replace RT with R2-G2 | Bayesian MLP · MNIST | test accuracy: 98.00 ± 0.11 -> 98.00 ± 0.08 | No clear effect |
| replace RT with R2-G2 | Bayesian CNN · CIFAR-10 | test log-likelihood: -3.88 ± 0.02 -> -3.87 ± 0.01 | Helped |
| replace RT with R2-G2 | Bayesian CNN · CIFAR-10 | test accuracy: 67.98 ± 0.88 -> 67.97 ± 0.66 | No clear effect |
| replace RT with R2-G2 | 2-layer VAE · MNIST | test ELBO: -107.64 ± 8.46 -> -106.85 ± 5.00 | Helped |
| replace RT with R2-G2 | 2-layer VAE · Omniglot | test ELBO: -131.48 ± 1.66 -> -129.80 ± 0.74 | Helped |
| replace RT with R2-G2 | 2-layer VAE · Fashion-MNIST | test ELBO: -240.59 ± 0.93 -> -240.23 ± 0.64 | Helped |
| replace RT with R2-G2 | 3-layer VAE · MNIST | test ELBO: -111.50 ± 5.40 -> -102.45 ± 3.73 | Helped |
| replace RT with R2-G2 | 3-layer VAE · Omniglot | test ELBO: -136.12 ± 2.70 -> -134.95 ± 2.26 | Helped |
| replace RT with R2-G2 | 3-layer VAE · Fashion-MNIST | test ELBO: -240.89 ± 0.97 -> -240.15 ± 0.86 | Helped |
| use R2-G2 instead of RT in a one-layer VAE (ablation) | VAE · MNIST | test ELBO: -94.22 ± 0.24 -> -94.39 ± 0.42 | Hurt |
| use R2-G2 instead of RT in a one-layer VAE (ablation) | VAE · Omniglot | test ELBO: -117.64 ± 2.12 -> -117.61 ± 2.12 | No clear effect |
Surprising findings
- Single-layer VAEs showed no gain from R2-G2 (RT slightly better on MNIST), which the authors attribute to reduced gradient variance hindering learning of stable representations.
- R2-G2 is roughly 3-4x slower in steps per second than RT for the MLP and CNN in the reported timings, with the gap depending on layer size.
Limitations
- Computational cost of solving the conjugate gradient system grows with the pre-activation dimension, O(m^2 n) extra for BNNs.
- Theory and method restricted to diagonal covariance V; non-diagonal case left to future work.
- Gains are sometimes modest and rely on initial training; the limited single-layer VAE results show no benefit.
Experimental setup
Datasets: MNIST CIFAR-10 Omniglot Fashion-MNIST
Benchmarks: MNIST CIFAR-10 MNIST VAE Omniglot Fashion-MNIST
Metrics: test log-likelihood test accuracy test variational lower bound gradient variance
Baselines: Reparameterisation (RT) estimator Local reparameterisation (LRT) estimator
Builds on: reparameterisation trick local reparameterisation trick Rao-Blackwellisation Adam
Closest papers by meaning
- The Generalized Reparameterization Gradient (NeurIPS 2016 · similarity 0.90)
- Reparameterization Gradient for Non-differentiable Models (NeurIPS 2018 · similarity 0.89)
- Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator (ICLR 2021 · similarity 0.88)
- Approximation Based Variance Reduction for Reparameterization Gradients (NeurIPS 2020 · similarity 0.87)
- Reducing Reparameterization Gradient Variance (NeurIPS 2017 · similarity 0.87)
- Reparameterization invariance in approximate Bayesian inference (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.