22
$\begingroup$

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: enter image description here

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

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 a small weight (0.001), I observe the following loss behavior: enter image description here 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:

enter image description here

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, interpolative latent space while producing reasonable reconstructions?

Improve this question
$\endgroup$

2 Answers 2

Reset to default
20
$\begingroup$

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.

Improve this answer
$\endgroup$
4
  • $\begingroup$ Interesting point, do you have a reference for that? $\endgroup$
    – dopexxx
    Commented Mar 2, 2019 at 11:20
  • 8
    $\begingroup$ @dopexxx Look up "KL annealing", e.g. arxiv.org/pdf/1511.06349.pdf $\endgroup$
    – jayelm
    Commented Apr 30, 2019 at 0:06
  • $\begingroup$ 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 $\endgroup$
    – dopexxx
    Commented Apr 30, 2019 at 0:08
  • 2
    $\begingroup$ @dopexxx just saying, Samy Bengio != Yoshua Bengio, just in. case you meant the latter. $\endgroup$
    – momo
    Commented May 9, 2021 at 2:49
7
$\begingroup$

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.

Improve this answer
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.