# Balancing Reconstruction vs KL Loss Variational Autoencoder

I am training a conditional variational autoencoder on a dataset of faces. When I set my KLL Loss equal to my Reconstruction loss term, my autoencoder seems unable to produce varied samples. I always get the same types of faces appearing:

These samples are terrible. However, when I decrease the weight of the KLL loss by 0.001, I get reasonable samples:

The problem is that the learned latent space is not smooth. If I try to perform a latent interpolation or generate a random sample, I get junk. When the KLL term has the small weight (0.001), I observe the following loss behavior: Notice that the VggLoss (the reconstruction term) decreases, while the KLLoss continues to increase.

I also tried increasing the dimensionality of the latent space, but this didn't work either.

Notice here, when the two loss terms are of equal weight, how the KLL term dominates, but doesn't allow the reconstruction loss to decrease:

This results in terrible reconstructions. Are there any suggestions on how to balance these two loss terms or any other possible things to try so that my autoencoder learns a smooth, interpolatable latent space, while producing reasonable reconstructions?

However, when I decrease the weight of the KLL loss by 0.001, I get reasonable samples: (...) The problem is that the learned latent space is not smooth.

Looks like overfitting. Remember that KL loss on the latent space sort of corresponds to regularization.

Are there any suggestions on how to balance these two loss terms or any other possible things to try so that my autoencoder learns a smooth, interpolatable latent space, while producing reasonable reconstructions?

I recently bumped into this paper: $\beta$-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework (it actually uses your dataset in one example).

From the paper ($\beta$ is the parameter you changed):

We introduce an adjustable hyperparameter $\beta$ that balances latent channel capacity and independence constraints with reconstruction accuracy (...) $\beta$-VAE is stable to train, makes few assumptions about the data and relies on tuning a single hyperparameter $\beta$, which can be directly optimised through a hyperparameter search using weakly labelled data or through heuristic visual inspection for purely unsupervised data.

Little late to the party here and you're probably way past this, but it's well documented you have to "warm up" the KL loss term by starting at zero and training a bit on just reconstruction loss before introducing the KL loss or results will not be good. It's unclear from your post if you did this, but it's a classic example of how touchy training these can be -- sometimes I wonder how appropriate they are on cross validated since it's very strongly trial and error, along with a bit of pixie dust and rainbows.

• Interesting point, do you have a reference for that? – dopexxx Mar 2 at 11:20
• @dopexxx Look up "KL annealing", e.g. arxiv.org/pdf/1511.06349.pdf – jayelm Apr 30 at 0:06
• Thanks! Who else than Bengio introduced that :D Actually I found the paper at some point later but had forgotten about this thread. It's a great reference – dopexxx Apr 30 at 0:08