1
$\begingroup$

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?

$\endgroup$

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.