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.
