Variational Auto-Encoder Loss function (keras)

Viewed 2683

I'm building a Variational auto-encoder using Keras. I was largely inspired by @Fchollet example : https://github.com/fchollet/keras/blob/master/examples/variational_autoencoder.py

but I'm working with continuous data. Instead of having to rank numbers like in MNIST, my output is a duration. In that regard, I changed the loss function from binary_crossentropy to mean_squared_error. I'm mostly wondering about the second term, the KL divergence. Is it supposed to work well with continuous data ? I can't wrap my head around it. To me, it's supposed to put the similar data tightly put together in the latent space. For example in cas of the MNIST data, put every "1" together in the latent space, every "2" together etc... Since I'm working with continuous data, how does it work ? Is there a better lost function to use in my case ?

Here's the lost function :

def vae_loss(x, x_decoded_mean):
    xent_loss = original_dim *metrics.mean_squared_error(x, x_decoded_mean)
    kl_loss = - 0.5 * K.mean(1 + z_log_sigma - K.square(z_mean) - K.exp(z_log_sigma), axis=-1)
    return K.mean(xent_loss + kl_loss)
vae.compile(optimizer='Adam', loss=vae_loss)

And here's a representation in the 3D latent space . enter image description here

As you can see, some similar data is put together as desired. Here's what happens when I increase the coefficient of the kl_loss function to "-100000" instead of "-0.5" : enter image description here

I thought I would have something going from blue to red in almost a linear way. Instead, I get a cluster of all the data in a chaotic fashion.

Can you guys please help me out ? Thank you !

1 Answers

In my opinion, this is because you increased the importance of the KL loss by increasing its coefficient. In other word, the loss function 'take care' of the KL term a lot more. Remember that the KL loss is used to 'fetch' the posterior distribution with the prior, N(0,1). So, if you 'expand' this behavior a lot (with a big coefficient), every data's distribution will converges toward this distribution ==> all data's posterior with mean 0 & data pushed together toward 0 ==> your latent representation :)

Related