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". Its use in variational inference was proposed in 2013.
Mathematics.
Let formula_1 be a random variable with distribution formula_2, where formula_3 is a vector containing the parameters of the distribution.
REINFORCE estimator.
Consider an objective function of the form:formula_4Without the reparameterization trick, estimating the gradient formula_5 can be challenging, because the parameter appears in the random variable itself. In more detail, we have to statistically estimate:formula_6The REINFORCE estimator, widely used in reinforcement learning and especially policy gradient, uses the following equality:formula_7This allows the gradient to be estimated:formula_8The REINFORCE estimator has high variance, and many methods were developed to reduce its variance.
Reparameterization estimator.
The reparameterization trick expresses formula_1 as:formula_10Here, formula_11 is a deterministic function parameterized by formula_3, and formula_13 is a noise variable drawn from a fixed distribution formula_14. This gives:formula_15Now, the gradient can be estimated as:formula_16
Examples.
For some common distributions, the reparameterization trick takes specific forms:
Normal distribution: For formula_17, we can use:formula_18
Exponential distribution: For formula_19, we can use:formula_20Discrete distribution can be reparameterized by the Gumbel distribution (Gumbel-softmax trick or "concrete distribution").
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:
formula_21
where formula_22 is the encoder (recognition model), formula_23 is the decoder (generative model), and formula_24 is the prior distribution over latent variables. The gradient of ELBO with respect to formula_25 is simplyformula_26but the gradient with respect to formula_3 requires the trick. Express the sampling operation formula_28 as:formula_29where formula_30 and formula_31 are the outputs of the encoder network, and formula_32 denotes element-wise multiplication. Then we haveformula_33where formula_34. This allows us to estimate the gradient using Monte Carlo sampling:formula_35where formula_36 and formula_37 for formula_38.
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:formula_39Using the reparameterization trick, we can estimate the gradient of this objective with respect to formula_3:formula_41
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:formula_42where formula_43 is the weight matrix, formula_44 is the input, and formula_45 are the (fixed) dropout rates.
More generally, other distributions can be used than the Bernoulli distribution, such as the gaussian noise:formula_46where formula_47 and formula_48, with formula_49 and formula_50 being the mean and variance of the formula_51-th output neuron. The reparameterization trick can be applied to all such cases, resulting in the "variational dropout" method.