Reparameterization trick explained

The reparameterization trick (aka "reparameterization gradient estimator") is a technique used in statistical machine learning, particularly in variational inference, variational autoencoders, and stochastic optimization. It allows for the efficient computation of gradients through random variables, enabling the optimization of parametric probability models using stochastic gradient descent, and the variance reduction of estimators.

It was developed in the 1980s in operations research, under the name of "pathwise gradients", or "stochastic gradients".^{[1] [2]} Its use in variational inference was proposed in 2013.^[3]

Mathematics

Let

be a random variable with distribution , where is a vector containing the parameters of the distribution.

REINFORCE estimator

Consider an objective function of the form:$$L(\backslash phi)\; =\; \backslash mathbb\_[f(z)]$$Without the reparameterization trick, estimating the gradient

can be challenging, because the parameter appears in the random variable itself. In more detail, we have to statistically estimate:$$\backslash nabla\_\backslash phi\; L(\backslash phi)\; =\; \backslash nabla\_\backslash phi\; \backslash int\; dz\; \backslash ;\; q\_\backslash phi(z)\; f(z)$$The REINFORCE estimator, widely used in reinforcement learning and especially policy gradient,^[4] uses the following equality:$$\backslash nabla\_\backslash phi\; L(\backslash phi)\; =\; \backslash int\; dz\; \backslash ;\; q\_\backslash phi(z)\; \backslash nabla\_\backslash phi(\backslash ln\; q\_\backslash phi(z))\; f(z)\; =\; \backslash mathbb\_[\backslash nabla\_\backslash phi(\backslash ln\; q\_\backslash phi(z))\; f(z)]$$This allows the gradient to be estimated:$$\backslash nabla\_\backslash phi\; L(\backslash phi)\; \backslash approx\; \backslash frac\; 1N\; \backslash sum\_^N\; \backslash nabla\_\backslash phi(\backslash ln\; q\_\backslash phi(z\_i))\; f(z\_i)$$The REINFORCE estimator has high variance, and many methods were developed to reduce its variance.^[5]

Reparameterization estimator

The reparameterization trick expresses

as:$$z\; =\; g\_\backslash phi(\backslash epsilon),\; \backslash quad\; \backslash epsilon\; \backslash sim\; p(\backslash epsilon)$$Here, is a deterministic function parameterized by , and is a noise variable drawn from a fixed distribution . This gives:$$L(\backslash phi)\; =\; \backslash mathbb\_[f(g\_\backslash phi(\backslash epsilon))]$$Now, the gradient can be estimated as:$$\backslash nabla\_\backslash phi\; L(\backslash phi)\; =\; \backslash mathbb\_[\backslash nabla\_\backslash phi\; f(g\_\backslash phi(\backslash epsilon))]\backslash approx\; \backslash frac\; 1N\; \backslash sum\_^N\; \backslash nabla\_\backslash phi\; f(g\_\backslash phi(\backslash epsilon\_i))$$

Examples

For some common distributions, the reparameterization trick takes specific forms:

Normal distribution

For

, we can use:$$z\; =\; \backslash mu\; +\; \backslash sigma\; \backslash epsilon,\; \backslash quad\; \backslash epsilon\; \backslash sim\; \backslash mathcal(0,\; 1)$$

Exponential distribution

For

, we can use:$$z\; =\; -\backslash frac\; \backslash log(\backslash epsilon),\; \backslash quad\; \backslash epsilon\; \backslash sim\; \backslash text(0,\; 1)$$Discrete distribution can be reparameterized by the Gumbel distribution (Gumbel-softmax trick or "concrete distribution").^[6]

In general, any distribution that is differentiable with respect to its parameters can be reparameterized by inverting the multivariable CDF function, then apply the implicit method. See for an exposition and application to the Gamma Beta, Dirichlet, and von Mises distributions.

Applications

Variational autoencoder

In Variational Autoencoders (VAEs), the VAE objective function, known as the Evidence Lower Bound (ELBO), is given by:

$$\backslash text(\backslash phi,\; \backslash theta)\; =\; \backslash mathbb\_[\backslash log\; p\_\backslash theta(x|z)]\; -\; D\_(q\_\backslash phi(z|x)\; ||\; p(z))$$

where

is the encoder (recognition model), is the decoder (generative model), and is the prior distribution over latent variables. The gradient of ELBO with respect to is simply$$\backslash mathbb\_[\backslash nabla\_\backslash theta\; \backslash log\; p\_\backslash theta(x|z)]\; \backslash approx\; \backslash frac\; \backslash sum\_^L\; \backslash nabla\_\backslash theta\; \backslash log\; p\_\backslash theta(x|z\_l)$$but the gradient with respect to requires the trick. Express the sampling operation as:$$z\; =\; \backslash mu\_\backslash phi(x)\; +\; \backslash sigma\_\backslash phi(x)\; \backslash odot\; \backslash epsilon,\; \backslash quad\; \backslash epsilon\; \backslash sim\; \backslash mathcal(0,\; I)$$where and are the outputs of the encoder network, and denotes element-wise multiplication. Then we have$$\backslash nabla\_\backslash phi\; \backslash text(\backslash phi,\; \backslash theta)\; =\; \backslash mathbb\_[\backslash nabla\_\backslash phi\; \backslash log\; p\_\backslash theta(x|z)\; +\; \backslash nabla\_\backslash phi\; \backslash log\; q\_\backslash phi(z|x)\; -\; \backslash nabla\_\backslash phi\; \backslash log\; p(z)]$$where . This allows us to estimate the gradient using Monte Carlo sampling:$$\backslash nabla\_\backslash phi\; \backslash text(\backslash phi,\; \backslash theta)\; \backslash approx\; \backslash frac\; \backslash sum\_^L\; [\backslash nabla\_\backslash phi\; \backslash log\; p\_\backslash theta(x|z\_l)\; +\; \backslash nabla\_\backslash phi\; \backslash log\; q\_\backslash phi(z\_l|x)\; -\; \backslash nabla\_\backslash phi\; \backslash log\; p(z\_l)]$$where and for .

This formulation enables backpropagation through the sampling process, allowing for end-to-end training of the VAE model using stochastic gradient descent or its variants.

Variational inference

More generally, the trick allows using stochastic gradient descent for variational inference. Let the variational objective (ELBO) be of the form:$$\backslash text(\backslash phi)\; =\; \backslash mathbb\_[\backslash log\; p(x,\; z)\; -\; \backslash log\; q\_\backslash phi(z)]$$Using the reparameterization trick, we can estimate the gradient of this objective with respect to

:$$\backslash nabla\_\backslash phi\; \backslash text(\backslash phi)\; \backslash approx\; \backslash frac\; \backslash sum\_^L\; \backslash nabla\_\backslash phi\; [\backslash log\; p(x,\; g\_\backslash phi(\backslash epsilon\_l))\; -\; \backslash log\; q\_\backslash phi(g\_\backslash phi(\backslash epsilon\_l))],\; \backslash quad\; \backslash epsilon\_l\; \backslash sim\; p(\backslash epsilon)$$

Dropout

The reparameterization trick has been applied to reduce the variance in dropout, a regularization technique in neural networks. The original dropout can be reparameterized with Bernoulli distributions:$$y\; =\; (W\; \backslash odot\; \backslash epsilon)\; x,\; \backslash quad\; \backslash epsilon\_\; \backslash sim\; \backslash text(\backslash alpha\_)$$where

is the weight matrix, is the input, and are the (fixed) dropout rates.

More generally, other distributions can be used than the Bernoulli distribution, such as the gaussian noise:$$y\_i\; =\; \backslash mu\_i\; +\; \backslash sigma\_i\; \backslash odot\; \backslash epsilon\_i,\; \backslash quad\; \backslash epsilon\_i\; \backslash sim\; \backslash mathcal(0,\; I)$$where

and , with and being the mean and variance of the -th output neuron. The reparameterization trick can be applied to all such cases, resulting in the variational dropout method.^[7]

See also

• Variational autoencoder
• Stochastic gradient descent
• Variational inference

Further reading

• Ruiz . Francisco R. . AUEB . Titsias RC . Blei . David . 2016 . The Generalized Reparameterization Gradient . Advances in Neural Information Processing Systems . 29 . 1610.02287 . September 23, 2024.
• Zhang . Cheng . Butepage . Judith . Kjellstrom . Hedvig . Mandt . Stephan . 2019-08-01 . Advances in Variational Inference . IEEE Transactions on Pattern Analysis and Machine Intelligence . 41 . 8 . 2008–2026 . 10.1109/TPAMI.2018.2889774 . 30596568 . 0162-8828. 1711.05597 .
• Web site: Mohamed . Shakir . October 29, 2015 . Machine Learning Trick of the Day (4): Reparameterisation Tricks . September 23, 2024 . The Spectator . .

Notes and References

1. Figurnov . Mikhail . Mohamed . Shakir . Mnih . Andriy . 2018 . Implicit Reparameterization Gradients . Advances in Neural Information Processing Systems . Curran Associates, Inc. . 31.
2. Fu, Michael C. "Gradient estimation." Handbooks in operations research and management science 13 (2006): 575-616.
3. Kingma . Diederik P. . Auto-Encoding Variational Bayes . 2022-12-10 . 1312.6114 . Welling . Max. stat.ML .
4. Williams . Ronald J. . 1992-05-01 . Simple statistical gradient-following algorithms for connectionist reinforcement learning . Machine Learning . en . 8 . 3 . 229–256 . 10.1007/BF00992696 . 1573-0565.
5. Greensmith . Evan . Bartlett . Peter L. . Baxter . Jonathan . 2004 . Variance Reduction Techniques for Gradient Estimates in Reinforcement Learning . Journal of Machine Learning Research . 5 . Nov . 1471–1530 . 1533-7928.
6. Maddison . Chris J. . The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables . 2017-03-05 . 1611.00712 . Mnih . Andriy . Teh . Yee Whye. cs.LG .
7. Kingma . Durk P . Salimans . Tim . Welling . Max . 2015 . Variational Dropout and the Local Reparameterization Trick . Advances in Neural Information Processing Systems . 28. 1506.02557 .