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?
