Constructing Discrete Probability Distributions from the Perspective of Reparameterization

从重参数的角度看离散概率分布的构建 苏剑林 · 2022-05-25 · Mathematics probability reparameterization

Generally speaking, the output of a neural network is unconstrained, i.e., it takes values in $\mathbb{R}$. To obtain a constrained output, one usually adds an activation function. For instance, if we want the output to be a probability distribution representing the probability of each category, we typically append a Softmax as the final activation function. A natural follow-up question then arises: besides Softmax, are there other operations that can produce a probability distribution?

In Random Musings on Reparameterization: From the Normal Distribution to Gumbel Softmax, we introduced the reparameterization trick behind Softmax. In this post, we will reverse that process: we first define a reparameterization operation, and then work backward to derive the corresponding probability distribution, giving us a new perspective for understanding the construction of probability distributions.

Problem Setup

Suppose the output vector of the model is $\boldsymbol{\mu}=[\mu_1,\cdots,\mu_n]\in\mathbb{R}^n$; without loss of generality, assume $\mu_i$ are pairwise distinct. We hope to use some transformation $\mathcal{T}$ to convert $\boldsymbol{\mu}$ into an $n$-way probability distribution $\boldsymbol{p}=[p_1,\cdots,p_n]$, while preserving certain properties. For instance, the most basic requirements are:

\begin{equation}{\color{red}1.}\,p_i\geq 0 \qquad {\color{red}2.}\,\sum_i p_i = 1 \qquad {\color{red}3.}\,p_i \geq p_j \Leftrightarrow \mu_i \geq \mu_j\end{equation}more

Of course, these requirements are quite mild — as long as $f$ is a monotonic function of $\mathbb{R}\mapsto\mathbb{R}^+$ (for Softmax we have $f(x)=e^x$), the transformation

\begin{equation}p_i = \frac{f(\mu_i)}{\sum\limits_j f(\mu_j)}\end{equation}

satisfies all of the above. Next, let's add a slightly less trivial condition:

\begin{equation}{\color{red}4.}\, \mathcal{T}(\boldsymbol{\mu}) = \mathcal{T}(\boldsymbol{\mu} + c\boldsymbol{1})\quad (\forall c \in \mathbb{R})\end{equation}

where $\boldsymbol{1}$ denotes the all-ones vector, and $c$ is an arbitrary constant. That is, adding the same constant to every component of $\boldsymbol{\mu}$ should leave the result of the transformation unchanged. We impose this condition because adding a constant to every component leaves $\mathop{\text{argmax}}$ unchanged, and it would be nice if $\mathcal{T}$ retained this same property as much as possible. It's easy to check that Softmax satisfies this condition; however, aside from Softmax, it seems hard to think of any other transformation that does.

Perturbation by Noise

Interestingly, we can construct such a transformation using the inverse of the reparameterization trick! Suppose $\boldsymbol{\varepsilon}=[\varepsilon_1,\cdots,\varepsilon_n]$ is a vector obtained by independently sampling from a distribution $p(\varepsilon)$ repeatedly $n$ times. Since $\boldsymbol{\varepsilon}$ is random, $\mathop{\text{argmax}}(\boldsymbol{\mu}+\boldsymbol{\varepsilon})$ is generally random as well, so we can define the transformation $\mathcal{T}$ via

\begin{equation}p_i = P[\mathop{\text{argmax}}(\boldsymbol{\mu}+\boldsymbol{\varepsilon})=i]\end{equation}

Since $\boldsymbol{\varepsilon}$ are i.i.d., and the whole definition only depends on $\mathop{\text{argmax}}(\boldsymbol{\mu}+\boldsymbol{\varepsilon})$ — that is, it only involves the relative magnitudes of each component — the transformation so defined necessarily satisfies the four conditions listed above.

We can also directly work out the form of $p_i$ to check which properties it satisfies. Specifically, $\mathop{\text{argmax}}(\boldsymbol{\mu}+\boldsymbol{\varepsilon})=i$ means that

\begin{equation}\mu_i + \varepsilon_i > \mu_j + \varepsilon_j\quad (\forall j\neq i)\end{equation}

that is, $\mu_i - \mu_j + \varepsilon_i > \varepsilon_j$. Clearly, the larger $\mu_i$ is, the more likely this equation holds, i.e., the larger $\mu_i$ is, the larger the corresponding $p_i$ — this is exactly condition $3$. Concretely, fixing $\varepsilon_i$, the probability that this condition holds is

\begin{equation}\int_{-\infty}^{\mu_i - \mu_j + \varepsilon_i} p(\varepsilon_j)d\varepsilon_j = \Phi(\mu_i - \mu_j + \varepsilon_i)\end{equation}

Here $\Phi$ is the cumulative distribution function (CDF) of $p(\varepsilon)$. Since each $\varepsilon_j$ is i.i.d., we can multiply these probabilities together directly:

\begin{equation}\prod_{j\neq i} \Phi(\mu_i - \mu_j + \varepsilon_i)\end{equation}

This is the probability of $\mathop{\text{argmax}}(\boldsymbol{\mu}+\boldsymbol{\varepsilon})=i$ given a fixed $\varepsilon_i$. Finally, we just need to average over $\varepsilon_i$ to obtain $p_i$:

\begin{equation}p_i = \int_{-\infty}^{\infty} p(\varepsilon_i)\left[\prod_{j\neq i} \Phi(\mu_i - \mu_j + \varepsilon_i)\right]d\varepsilon_i \label{eq:pi}\end{equation}

From the expression for $p_i$, we can see that it only depends on the relative values $\mu_i - \mu_j$, and therefore it clearly satisfies condition $4$ from our definition.

Revisiting Familiar Ground

Comparing this with the discussion of Gumbel-Max in Random Musings on Reparameterization: From the Normal Distribution to Gumbel Softmax, we can see that the derivation above runs exactly opposite to reparameterization: here we first define the reparameterization method, and then derive the corresponding probability distribution backward.

Now let's re-examine our earlier result: when the noise distribution is chosen to be the Gumbel distribution, does equation $\eqref{eq:pi}$ reduce to the standard Softmax operation? Gumbel noise $u\sim U[0,1]$ is obtained by transforming $\varepsilon = -\log(-\log u)$, and since the distribution of $u$ is exactly $U[0,1]$, solving for $u=e^{-e^{-\varepsilon}}$ gives precisely the CDF of the Gumbel distribution, i.e., $\Phi(\varepsilon)=e^{-e^{-\varepsilon}}$, while $p(\varepsilon)$ is the derivative of $\Phi(\varepsilon)$, i.e., $p(\varepsilon)=\Phi'(\varepsilon)=e^{-\varepsilon-e^{-\varepsilon}}$.

Substituting these results into equation $\eqref{eq:pi}$ gives

\begin{equation}\begin{aligned} p_i =&\, \int_{-\infty}^{\infty} e^{-\varepsilon_i-e^{-\varepsilon_i}} e^{-\sum\limits_{j\neq i}e^{-\varepsilon_i + \mu_j - \mu_i}} d\varepsilon_i \\ =&\, \int_{-\infty}^0 e^{-e^{-\varepsilon_i}\left(1+\sum\limits_{j\neq i}e^{\mu_j - \mu_i}\right)} d(-e^{-\varepsilon_i}) \\ =&\, \int_{-\infty}^0 e^{t\left(1+\sum\limits_{j\neq i}e^{\mu_j - \mu_i}\right)} dt\\ =&\, \frac{1}{1+\sum\limits_{j\neq i}e^{\mu_j - \mu_i}} = \frac{e^{\mu_i}}{\sum\limits_j e^{\mu_j }} \end{aligned}\end{equation}

which is exactly Softmax. This again confirms the correspondence between Gumbel-Max and Softmax.

Numerical Computation

Finding an analytical solution as clean as Softmax's — as in the Gumbel-distribution case — is extremely rare; at least the author hasn't been able to find a second example so far. So in most cases, we can only rely on numerical methods to approximate equation $\eqref{eq:pi}$. Since $p(\varepsilon)=\Phi'(\varepsilon)$, we can directly manipulate the differential to get:

\begin{equation}p_i = \int_0^1 \left[\prod_{j\neq i} \Phi(\mu_i - \mu_j + \varepsilon_i)\right]d\Phi(\varepsilon_i)\end{equation}

Let $t=\Phi(\varepsilon_i)$, so that

\begin{equation}\begin{aligned} p_i =&\, \int_0^1 \left[\prod_{j\neq i} \Phi(\mu_i - \mu_j + \Phi^{-1}(t))\right]dt \\ \approx&\, \frac{1}{K}\sum_{k=1}^K\prod_{j\neq i} \Phi\left(\mu_i - \mu_j + \Phi^{-1}\left(\frac{k}{K+1}\right)\right) \end{aligned}\end{equation}

where $\Phi^{-1}$ is the inverse function of $\Phi$, also known in probability theory as the quantile function (or percent-point function).

From the equation above, we see that as long as we know the analytical form of $\Phi$, we can approximately compute $p_i$. Note that we don't actually need an analytical form for $\Phi^{-1}$, because the results at the sampled points $\Phi^{-1}\left(\frac{k}{K+1}\right)$ can be precomputed ahead of time using other numerical methods.

Take the standard normal distribution as an example: $\Phi(x)=\frac{1}{2} \left(1+\text{erf}\left(\frac{x}{\sqrt{2}}\right)\right)$, and mainstream deep learning frameworks all come with a built-in $\text{erf}$ function, so computing $\Phi(x)$ poses no problem; as for $\Phi^{-1}\left(\frac{k}{K+1}\right)$, we can precompute it using scipy.stats.norm.ppf. So when $\boldsymbol{\varepsilon}$ is sampled from the standard normal distribution, computing $p_i$ is entirely feasible in mainstream deep learning frameworks.

Summary

This post generalizes Softmax from the perspective of reparameterization, yielding a family of probability normalization methods that share similar properties.