Reparametrization trick

In machine learning, the reparametrization trick is a method to construct unbiased estimators of gradients, typically used for performing gradient descent on expectations.

It was introduced independently to variational inference by Kingma & Welling (2014), Rezende et al. (2014), and Titsias & Lazaro-Gredilla (2014).

Mathematical principle

A problem commonly encountered in machine learning is estimating gradient-of-expectation. The problem is as follows:

Given a family of distributions

$\{p_{\theta }{\}}_{\theta \in \mathrm{\Theta }}$

, where each

$\theta$

defines a probability measure

$p_{\theta }(x)dx$

on a base space

${𝒳}$

.

Given any nice (nice enough to allow differentiating under the integral sign) function

$f_{\theta }:{𝒳}\to \mathbb{ℝ}$

, the gradient-of-expectation is:

${\mathrm{\nabla }}_{\theta }{E}_{x\sim p_{\theta }}[f_{\theta }(x)]=\int {\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)p_{\theta }(x)dx+\int f_{\theta }(x){\mathrm{\nabla }}_{\theta }{p}_{\theta }(x)dx$

(main)

The key problem is to estimate the gradient-of-expectation efficiently. There are several possible solutions.

Analytic

If the expectation has closed-form solution, then its gradient also has closed-form solution. This is essentially the only method used by statisticians before large-scale computing was available.

There are several problems with this. One, the closed-form solution rarely exists. Two, the closed-form solution might be very expensive to compute, especially when $\mathrm{\Theta }$ has a large dimension (curse of dimensionality). The typical solution for such intractable problems is Monte Carlo method.

Monte Carlo method

The naive method to estimating Equation (main) is to sample many $x_i\sim p_{\theta }$, estimate ${\mathbb{𝔼}}_{x\sim p_{\theta }}[f_{\theta }(x)]$, then vary $\theta$ slightly to $\theta +\delta {\theta }_j$, estimate ${\mathbb{𝔼}}_{x\sim p_{\theta +\delta {\theta }_j}}[f_{\theta +\delta {\theta }_j}(x)]$, and so on, and finally fitting a linear operator $L$ such that$${\mathbb{𝔼}}_{x\sim p_{\theta +\delta {\theta }_j}}[f_{\theta +\delta {\theta }_j}(x)]-{\mathbb{𝔼}}_{x\sim p_{\theta }}[f_{\theta }(x)]\approx L(\delta {\theta }_j)\mathrm{\forall }j$$This is highly inefficient. To save time, one should instead compute the gradient exactly where it is tractable, instead of estimating the gradient everywhere. In most commonly used examples, while ${\mathrm{\nabla }}_{\theta }{E}_{x\sim p_{\theta }}[f_{\theta }(x)]$ is intractable, ${\mathrm{\nabla }}_{\theta }{f}_{\theta }(x),{\mathrm{\nabla }}_{\theta }{p}_{\theta }(x)$ are tractable (they are often designed by practitioners to be tractable).

As such, we consider the following expansion of Equation (main):

${\mathrm{\nabla }}_{\theta }{E}_{x\sim p_{\theta }}[f_{\theta }(x)]=\int {\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)p_{\theta }(x)dx+\int f_{\theta }(x){\mathrm{\nabla }}_{\theta }{p}_{\theta }(x)dx$

(Monte Carlo)

Instead of estimating the integrals, then estimating the gradient, now we compute the gradients exactly, then estimate the integrals. The integrals are still intractable in general, but can be done by Monte Carlo integration. To perform Monte Carlo integration, an integral

$\underset{x\in {𝒳}}{\int }g(x)dx$

must have a probability distribution to sample

$x$

from. If we use

$p_{\theta }$

for Monte Carlo integration, we obtain

$$\int \frac{{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)p_{\theta }(x)+f_{\theta }(x){\mathrm{\nabla }}_{\theta }{p}_{\theta }(x)}{p_{\theta }(x)}{p}_{\theta }(x)dx=E_{x\sim p_{\theta }}[{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)+f_{\theta }(x){\mathrm{\nabla }}_{\theta }\ln p_{\theta }(x)]$$

and thus

${\mathrm{\nabla }}_{\theta }{E}_{x\sim p_{\theta }}[f_{\theta }(x)]=E_{x\sim p_{\theta }}[{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)+f_{\theta }(x){\mathrm{\nabla }}_{\theta }\ln p_{\theta }(x)]$

(REINFORCE)

This is the equation used in policy gradient methods, to be detailed below.

The reparametrization trick

The main issue in estimating Equation (main) is an "entanglement" between the distribution

$x\sim p_{\theta }$

and the expectation of

$f_{\theta }(x)$

to be estimated. The entanglement comes to the fore when we vary

$\theta$

. The reparametrization trick pushes all dependence on

$\theta$

into a deterministic function, and then perform the obvious

$${\mathrm{\nabla }}_{\theta }{\mathbb{𝔼}}_{x\sim p}[f_{\theta }(x)]={\mathbb{𝔼}}_{x\sim p}[{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)]$$

The prototypical example is the family of 1D normal distributions:

$\{N(\mu ,{\sigma }^2){\}}_{\mu \in \mathbb{ℝ},\sigma >0}$

. Given any

$\mu \in \mathbb{ℝ},\sigma >0$

, we can sample

$x\sim N(\mu ,{\sigma }^2)$

as

$x=f_{\mu ,\sigma }(x^{\prime })$

, with

$x^{\prime }\sim N(0,1)$

and

$f_{\mu ,\sigma }(x)=\mu +\sigma x$

.

Remark: The idea is similar to probability transforms such as the Box–Muller transform, where we have only one "seed" random number generator, and must construct other probability distributions by performing deterministic transforms on the random numbers generated by it.

In general, given a family of distribution

$\{p_{\theta }{\}}_{\theta \in \mathrm{\Theta }}$

, we can perform the reparametrization trick if there exists a seed distribution

$p$

, and a transform function

$g_{\theta }(x)$

, such that for any

$\theta \in \mathrm{\Theta }$

, we can sample from

$p_{\theta }$

by sampling

$x\sim p$

, then compute

$g_{\theta }(x)$

. That is,

$p_{\theta }=p\circ g_{\theta }^{-1}$

. With the seed distribution and the transform, we obtain the reparametrization trick equation:

${\mathrm{\nabla }}_{\theta }{E}_{x\sim p_{\theta }}[f_{\theta }(x)]={\mathrm{\nabla }}_{\theta }{E}_{x\sim p}[f_{\theta }(g_{\theta }(x))]=E_{x\sim p}[{\mathrm{\nabla }}_{\theta }{f}_{\theta }(g_{\theta }(x))]$

(reparametrization trick)

Motivation

Many problems in statistics and machine learning are of the form: find the "best" parameters. Usually, "best" is defined as "achieving minimal loss" or "maximum reward", and taking the gradient, if possible, is useful to optimization.

Parameter estimation in statistics

For example, consider the parameter estimation problem: Given data $x\in {𝒳}$ sampled from a distribution $p_{{\theta }_0}$, or at least from a distribution close enough to some $p_{{\theta }_0}$. The problem is to estimate ${\theta }_0$. As we will see below, this often reduces to solving one of the following problems:$$\underset{\theta }{\min }{\mathbb{𝔼}}_{x\sim p}[f_{\theta }(x)];\underset{\theta }{\min }{\mathbb{𝔼}}_{x\sim p_{\theta }}[f(x)]$$This is no loss of generality, because if we have a sequence of independently sampled data $x_1,...,x_n\in {𝒳}$, then we can simply consider them as one big data $x:=(x_1,...,x_n)\in {{𝒳}}^n$ with distribution $p_{{\theta }_0}^{\otimes n}$:$$p_{{\theta }_0}^{\otimes n}(x)dx=p_{{\theta }_0}(x_1)\cdots p_{{\theta }_0}(x_n)d{x}_1\cdots d{x}_n$$First approach (Blackwell 1951): find a distribution that can best mock-up the observed samples.$$\underset{\theta }{\min }{\mathbb{𝔼}}_{x^{\prime }\sim p_{\theta }}[d(x,x^{\prime })]$$where $d:{𝒳}\times {𝒳}\to [0,\mathrm{\infty })$ measures the difference between two points in ${𝒳}$. It can be a metric, or be more general.

Set $f(x^{\prime }):=d(x,x^{\prime })$, then this problem reduces to:$$\underset{\theta }{\min }{\mathbb{𝔼}}_{x^{\prime }\sim p_{\theta }}[f(x^{\prime })]$$

Second approach (frequentist estimation): define a loss function $L:\mathrm{\Theta }\times \mathrm{\Theta }\to [0,\mathrm{\infty })$, which measures the difference between two "points" (each point is a probability distribution, very big points indeed) in the space of distributions under consideration.

Then, minimize estimation risk (expected loss):$$\underset{\zeta \in \mathrm{Z}}{\min }{\mathbb{𝔼}}_{x\sim p_{{\theta }_0}}[L({\hat{\theta }}_{\zeta }(x),{\theta }_0)]$$where ${\hat{\theta }}_{\zeta }:{𝒳}\to \mathrm{\Theta }$ is an estimator, parametrized by $\zeta \in \mathrm{Z}$.

The above definition is not very interesting, since we could define the following "blind guess" estimator $\hat{\theta }(x)={\theta }_1$. It would happen to be exactly right if ${\theta }_1={\theta }_0$. Thus, we must additionally require the estimator to perform well on many different possible ${\theta }_0$.

There are many possible ways to formalize the idea of "perform well on many different ${\theta }_0$". A common version is the following:$$\underset{\zeta \in \mathrm{Z}}{\min }\underset{{\theta }_0\in \mathrm{\Theta }}{\max }{\mathbb{𝔼}}_{x\sim p_{{\theta }_0}}[L({\hat{\theta }}_{\zeta }(x),{\theta }_0)]$$Set $f_{\zeta }(x,{\theta }_0)=L({\hat{\theta }}_{\zeta }(x),{\theta }_0)$, then this problem reduces to:$$\underset{\zeta \in \mathrm{Z}}{\min }(\underset{{\theta }_0\in \mathrm{\Theta }}{\max }{\mathbb{𝔼}}_{x\sim p_{{\theta }_0}}[f_{\zeta }(x,{\theta }_0)])$$which is almost in the form of $\underset{\theta }{\min }{\mathbb{𝔼}}_{x\sim p}[f_{\theta }(x)]$, but not quite.

For example, when $\mathrm{\Theta }$ is a subset of ${\mathbb{ℝ}}^n$, $L(\theta ,{\theta }_0):=\Vert \theta -{\theta }_0{\Vert }^2$, and the set of estimators $\{{\hat{\theta }}_{\zeta }{\}}_{\zeta \in \mathrm{Z}}$ contains only unbiased estimators, then if the minimum-variance unbiased estimator exists, it is the solution to the above problem.

Third approach (Bayesian estimation): Imposing a maximum gives frequentist estimation a kind of "game-theoretic" flavor, since it is formally equivalent to a zero-sum game between a statistician and nature. The statistician proposes an estimator ${\hat{\theta }}_{\zeta }$, and nature replies with $p_{{\theta }_0}$.

However, nature is uninterested, as the statistician's choice of ${\hat{\theta }}_{\zeta }$ has no effect on $p_{{\theta }_0}$. Consequently, Bayesian estimation models this by imposing a prior distribution $\mu$ over $\mathrm{\Theta }$, and optimizing the following:$$\underset{\zeta \in \mathrm{Z}}{\min }{\mathbb{𝔼}}_{\theta \sim \mu }[{\mathbb{𝔼}}_{x\sim p_{\theta }}[L({\hat{\theta }}_{\zeta }(x),\theta )]]$$Set $f_{\zeta }(x,\theta )=L({\hat{\theta }}_{\zeta }(x),\theta )$, then this problem reduces to:$$\underset{\zeta \in \mathrm{Z}}{\min }{\mathbb{𝔼}}_{\theta \in \mathrm{\Theta },x\sim p_{\theta }}[f_{\zeta }(x,\theta )]$$which is in the form of $\underset{\theta }{\min }{\mathbb{𝔼}}_{x\sim p}[f_{\theta }(x)]$.

Reinforcement learning

In reinforcement learning, there is an "entanglement" between the distribution and the function.

To perform a gradient descent, one must estimate the gradient.

Main methods

There are many different ways to perform reparametrization trick, for diverse purposes.

Reparametrizing a distribution family

Given a family of distributions $\{p_{\theta }{\}}_{\theta }$, we can apply the reparametrization trick if we have a way to generate the family into a constant "seed random generator" and a family of parametrized deterministic functions.

For example, the family of normal distributions on ${\mathbb{ℝ}}^n$ is $\{N(\mu ,\mathrm{\Sigma }){\}}_{\mu \in {\mathbb{ℝ}}^n,\mathrm{\Sigma }⪰0}$. Given any $\mu \in {\mathbb{ℝ}}^n,\mathrm{\Sigma }⪰0$, we can sample $x\sim N(\mu ,\mathrm{\Sigma })$ as $x=f_{\mu ,\mathrm{\Sigma }}(x^{\prime })$, with $x^{\prime }\sim N(0,I_{n\times n})$ and $f_{\mu ,\mathrm{\Sigma }}(x)=\mu +Mx$, where $M{M}^T=\mathrm{\Sigma }$.

Since $\mathrm{\Sigma }$ is non-negative-definite, $M$ is guaranteed to exist by the spectral theorem, but it is not unique. It can be found by Cholesky decomposition, or singular value decomposition. Different choices have different theoretical and practical advantages.

Gumbel max tricks

The prototype of Gumbel tricks is the Gumbel-max trick, which allows one to create any categorical distribution using just a Gumbel distribution random number generator.

The Gumbel-max trick generalizes to other distributions:

Theorem — Gumbel

Template:Proof

This provides a proof that the Gumbel, Weibull, and Fréchet distributions are max-stable, which is one-half of the extreme value theorem.

Gumbel softmax method

The Gumbel-max trick allows sampling from the categorical distribution, but it cannot be used for training with gradient descent, because ${\mathrm{\nabla }}_x\arg \underset{i}{\max }(x_i+g_i)=0$. This is because $\arg \underset{i}{\max }(x_i+g_i)$ is "hard", that is, insensitive for small variations of $x$. Intuitively speaking, this can be interpreted as saying that the model simply predicts "category $i$ is most likely" without saying by how much it is the most likely category.

To create better gradients, the model should predict a distribution over the categories, and the standard method is the softmax function, creating the Gumbel softmax method:$$p_i:=\frac{e^{\beta (x_i+g_i)}}{\sum _j{e}^{\beta (x_j+g_j)}}$$After imposing a good loss function, such as the cross-entropy loss, one can train the method by standard gradient descent:$$Loss=-\ln p_i=-\ln \frac{e^{\beta (x_i+g_i)}}{\sum _j{e}^{\beta (x_j+g_j)}}$$where $i$ is the correct label.

Estimating bounds of partition functions

Inspired by the connection between information theory and statistical mechanics, energy-based models, such as the Boltzmann machine and the deep belief network, are statistical models defined by an energy function (or "potential function") and a temperature.

Consider a set of discrete random variables $X_1,...,X_n$, each $X_i$ being the state of a particle $i$. Let the system of $n$ particles interact, and let the energy of the entire system be $\varphi (x_1,...,x_n)$, when $X_1=x_1,...,X_n=x_n$. Then, when the system is in contact with a heat bath with temperature $T$, after reaching equilibrium, the distribution of the states of the system is the Boltzmann distribution:$$Pr(X_1=x_1,...,X_n=x_n|\beta )=\frac{e^{-\beta \varphi (x_1,...,x_n)}}{\sum _{x{\text{'}}_1,...,x{\text{'}}_n}{e}^{-\beta \varphi (x{\text{'}}_1,...,x{\text{'}}_n)}}$$where $\beta =\frac{1}{T}$ is the inverse temperature of the heat bath.

The normalizing constant $$Z(\beta )=\sum _{x{\text{'}}_1,...,x{\text{'}}_n}{e}^{-\beta \varphi (x{\text{'}}_1,...,x{\text{'}}_n)}$$is the partition function of the system. It depends on the temperature and the energy function.

In general, partition functions are intractable, making it important to estimate it in practice. There are a family of reparametrization tricks for accomplishing the estimation.

Upper bounds

Lower bounds

Applications

Variational autoencoder

In variational autoencoder,

Two similar methods

Related methods

The reparametrization trick is a general technique

REINFORCE

The policy gradient method in reinforcement learning, proposed in (Williams, 1992), uses the following expansion of Equation (main):

${\mathrm{\nabla }}_{\theta }{E}_{x\sim p_{\theta }}[f_{\theta }(x)]=E_{x\sim p_{\theta }}[{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)+f_{\theta }(x){\mathrm{\nabla }}_{\theta }\ln p_{\theta }(x)]$

(policy gradient)

It is also called the "likelihood ratio method" and the "score function method".

The policy gradient method has many variants, such as with function approximation, deterministic. An introduction is .

References

This article "Reparametrization trick" is from Wikipedia. The list of its authors can be seen in its historical and/or the page Edithistory:Reparametrization trick. Articles copied from Draft Namespace on Wikipedia could be seen on the Draft Namespace of Wikipedia and not main one.

---