Tensorflow Probability Implementation of Automatic Differentiation Variational Inference with Mixtures

In this paper, the authors suggest using the following loss instead of the traditional ELBO in order to train what basically is a Variational Autoencoder with a Gaussian Mixture Model instead of a single, normal distribution: $$\mathcal{L}_{SIWAE}^T(\phi)=\mathbb{E}_{\{z_{kt}\sim q_{k,\phi}(z|x)\}_{k=1,t=1}^{K,T}}\left[\log\frac{1}{T}\sum_{t=1}^T\sum_{k=1}^K\alpha_{k,\phi}(x)\frac{p(x|z_{k,t})r(z_{kt})}{q_\phi(z_{kt}|x)}\right]$$ They also provide the following code which is supposed to be a tensorflow probability implementation:

def siwae(prior, likelihood, posterior, x, T):
q = posterior(x)
z = q.components_dist.sample(T)
z = tf.transpose (z, perm=[2, 0, 1, 3])
loss_n = tf.math.reduce_logsumexp(
(−tf.math.log(T) + tf.math.log_softmax(mixture_dist.logits)[:, None, :]
+ prior.log_prior(z) + likelihood(z).log_prob(x) − q.log_prob(z)), axis=[0, 1])
return tf.math.reduce_mean(loss_n, axis=0)


However, it seems like this doesn't work at all so as someone with nearly no tensorflow knowledge I came up with the following:

def siwae(prior, likelihood, posterior, x, T):
q = posterior(x) # distribution over variables of shape (batch_size, 2)
z = q.components_distribution.sample(T)
z = tf.transpose(z, perm=[2, 0, 1, 3]) # shape (K, T, batch_size, encoded_size)
l1 = -tf.math.log(float(T)) # shape: (), log (1/T)
l2 = tf.math.log_softmax(tf.transpose(q.mixture_distribution.logits))[:, None , :] # shape (K, 1, batch_size), alpha
l3 = prior.log_prob(z) # shape (K, T, batch_size), r(z)
l4 = likelihood(tf.reshape(z, (K*T*x.shape[0], encoded_size)))
l4 = l4.log_prob(tf.repeat(x, repeats=K*T, axis=0)) # shape (K*T*batch_size, )
l4 = tf.reshape(l4, (K, T, x.shape[0])) # shape (K, T, batch_size), p(x|z)
l5 = -q.log_prob(z) # shape (K, T, batch_size), q(z|x)
loss_n = tf.math.reduce_logsumexp(l1 + l2 + l3 + l4 + l5, axis=[0, 1])
return tf.math.reduce_mean(loss_n, axis=0)


There are no errors when I try to use this as

siwae(prior, decoder, encoder, x_test[:100, ...], T)


but after a few training steps I get only nans. I really don't have any idea of this is an due to a wrong implementation or wrong usage of the loss - especially as I don't have much experience with tensorflow. So any help would be greatly appreciated. For a full, minimal example I created this colab.

• Hi, Jonas. Unfortunately, this AI stack is not focused on too technical and computational problems. You should try asking the same question on other Stacks, like: datascience.stackexchange.com Aug 25 at 13:46
• Hi Andre, thanks for the suggestion, I was a bit unsure on where to post and ended up here but in this case I'll try datascience. Aug 25 at 14:27
• Or would you suggest to even choose stackoverflow instead @AndreGoulart? My problem is that I'm not even sure myself whether the problem is just the tensorflow code (-> stackoverflow) or more on the side of getting the theory wrong. Aug 25 at 14:33