# Did I understand the usage of Gumbel-Softmax reparametrization correctly?

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?