KL divergence in VAE

[marianacuri]

I am trying to modify the loss function in the VAE example available here: https://github.com/rstudio/keras/blob/master/vignettes/examples/variational_autoencoder.R

In the example, the KL divergence is defined assuming independence among the latent variables (and standard normal for the prior), i.e.: kl_loss <- -0.5*k_mean(1 + z_log_var - k_square(z_mean) - k_exp(z_log_var), axis = -1L)

I would like to substitute this line in the code by another that specifies the KL divergence between two multivariate normal distributions (with non-zero covariances), i.e:

, where d is the number of latent variables,

I included one hidden layer to the endoder output, with the covariance parameters: num_cov <- latent_dim*(latent_dim)/2 z_cov <- layer_dense(h, num_cov, name='z_cov')

How should _klloss be rewritten in the code? And the sampling function to sample from a multivariate normal, instead from a standard normal?

[skeydan]

I suggest you look into tfprobability for tasks like this: https://rstudio.github.io/tfprobability/

On that page, under the heading "keras layers", there is a VAE example using a multivariate normal.

This post: https://medium.com/tensorflow/variational-autoencoders-with-tensorflow-probability-layers-d06c658931b7 has background explanations (at some time on our blog we'll have a post about using keras layers in tfprobability too :-)).