I am working on a deep learning model, which has a mixture of experts formulation like $\log p(y|x)=\log \sum_{z}p(y|z,x,\theta)p(z|x,\phi)$. So, each $p(y|z,x,\theta)$ is a deep learning classifier, whose parameters are $\theta$ and $p(z|x,\phi)$ defines a distribution on the expert components (which is also a deep learning model, with the parameters $\phi$.)

I am aiming to build the lower bound of $\log p(y|x)$ as $\sum_{z}\log p(y|z,x,\theta)p(z|x,\phi)$ and optimize this bound instead since $\log p(y|x)$ is not easy to optimize directly.

This new bound is an expectation with respect to a discrete variable $z$. At this point, I want to grasp exactly how to approach this problem and whether I understood the usage of Gumbel-Softmax relaxation correctly.

The derivative wrt $\theta$ is easy, it is $\sum_{z}\nabla_{\theta}\log p(y|z,x,\theta)p(z|x,\phi)$ and it can be approximated with Monte Carlo sampling from the dataset and each expert $z$.

The derivative of $\phi$ is $\sum_{z}\log p(y|z,x,\theta) \nabla_{\phi}p(z|x,\phi)$. This cannot be approximated since it is not an expectation. So we either need a REINFORCE-like estimator or a reparametrization trick. Since it is a discrete distribution, we need Gumbel-Softmax reparametrization.

I read the both papers on the Gumbel-Softmax reparametrization and I think I understood the mathematics behind it. But I am not sure how to implement it in my case. I want to be sure if the following would be correct:

We sample $g$ from the Gumbel distribution $G(g|0,1)$. Then the output probabilities from the network $\alpha = p(z|x,\phi)$ are taken and the surrogate variables are calculated: $\beta(x,g)_i = \dfrac{\exp \{ (\log \alpha_i + g_i)/\lambda \}}{\sum_j\exp \{ (\log \alpha_j + g_j)/\lambda \}}$. It is now $p(z|x,\phi) \approx \beta(x,g)$.

So, the lower bound will now be:

$$\int \log p(y|z =onehot(\beta(x,g)),x_i,\theta)\beta(x,g)G(g|0,1)dg$$. Onehot, since we will use the $z$ to which $\beta(x,g)$ will be closest on the probability simplex.

So the derivative $\nabla_{\phi}$ will be approximately now:

$$\nabla_{\phi} \approx \dfrac{1}{n}\sum_{i=1}^{n}\log p(y|z =onehot(\beta(x,g_i)),x,\theta)\nabla_{\phi}\beta(x,g_i)$$

Of course; the training will involve some tinkering with the annealing schedule of $\lambda$.

Is this understanding correct about the reparametrization trick and the Gumbel-Softmax version of it?


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.