07 November 2016
Notes on the Concrete distribution (continuous relaxations of discrete distributions) which was introduced by (Maddison, Mnih, & Teh, 2016) and (Jang, Gu, & Poole, 2016). Notes are based on the former.
Given a random variable with distribution parameterized by ; and a function . The goal is to approximate using the reparameterization trick (instead of the REINFORCE trick) when is discrete, i.e. when is countable.
The support of a -dimensional Concrete distribution is the simplex \begin{align} \Delta^{K - 1} = \left\{x \in \mathbb R^K: x_k \in [0, 1], \sum_{k = 1}^K x_k = 1\right\}. \end{align} A sample from the Concrete distribution—parameterized by location parameters , and temperature —can be generated as follows:
For :
The Concrete distribution is actually defined by its density (see equation (11) in (Maddison, Mnih, & Teh, 2016)). As is smaller, the approximation to the one-hot discrete random variable with unnormalized weights is more accurate in the sense of Proposition 1 in (Maddison, Mnih, & Teh, 2016).
Let . Let be distributed from , and the th output of the function be: \begin{align} \left[{g_{\phi}(Z)}\right]_k := \frac{\exp((\log \alpha_k + Z_k)/\lambda)}{\sum_{i = 1}^K \exp((\log \alpha_i + Z_i)/\lambda)}. \end{align}
Hence, is a Concrete random variable and we can write: \begin{align} \E_{X \sim \mathrm{Concrete}}[f(X)] &= \E_{Z \sim \prod_k \mathrm{Gumbel}(0, 1)}[f(g_{\phi}(Z))]. \end{align} Hence the derivatives of the objective will be: \begin{align} \frac{\partial}{\partial \phi} \E_{X \sim \mathrm{Concrete}}[f(X)] &= \frac{\partial}{\partial \phi} \E_{Z \sim \prod_k \mathrm{Gumbel}(0, 1)}[f(g_{\phi}(Z))] \\ &= \E_{Z \sim \prod_k \mathrm{Gumbel}(0, 1)}\left[\frac{\partial}{\partial \phi} \left(f(g_{\phi}(Z))\right)\right]. \end{align} This can be approximated by the Monte Carlo estimator: \begin{align} &\frac{1}{N} \sum_{n = 1}^N \frac{\partial}{\partial \phi} \left(f(g_{\phi}(Z^n))\right) \\ Z^n &:= (Z_1^n, \dotsc, Z_K^n) & n = 1, \dotsc, N \\ Z_k^n &\sim \mathrm{Gumbel}(0, 1) & k = 1, \dotsc, K, n = 1, \dotsc, N. \end{align}
References
@article{maddison2016concrete, title = {The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables}, author = {Maddison, Chris J. and Mnih, Andriy and Teh, Yee Whye}, journal = {arXiv preprint arXiv:1611.00712}, year = {2016} }
@article{jang2016categorical, title = {Categorical Reparameterization with Gumbel-Softmax}, author = {Jang, Eric and Gu, Shixiang and Poole, Ben}, journal = {arXiv preprint arXiv:1611.01144}, year = {2016} }
[back]