Tuan Anh Le

Concrete distribution

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.

Concrete distribution

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 :

  1. Generate a sample from a Gumbel distribution by setting where is a sample from .
  2. Set \begin{align} X_k \leftarrow \frac{\exp((\log \alpha_k + G_k)/\lambda)}{\sum_{i = 1}^K \exp((\log \alpha_i + G_i)/\lambda)}. \end{align}

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).

Reparameterization trick

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}


  1. Maddison, C. J., Mnih, A., & Teh, Y. W. (2016). The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. ArXiv Preprint ArXiv:1611.00712.
      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}
  2. Jang, E., Gu, S., & Poole, B. (2016). Categorical Reparameterization with Gumbel-Softmax. ArXiv Preprint ArXiv:1611.01144.
      title = {Categorical Reparameterization with Gumbel-Softmax},
      author = {Jang, Eric and Gu, Shixiang and Poole, Ben},
      journal = {arXiv preprint arXiv:1611.01144},
      year = {2016}