1
$\begingroup$

I am training a VAE on CelebA HQ (resized to 256x256). The training is going well, the reconstruction loss is decreasing and reconstructions are also meaningful. But, the problem is with KL divergence loss. The KL divergence loss first drops, then start to increase.

training logs file

I have attached link to training logs file.

$\endgroup$

1 Answer 1

1
$\begingroup$

That is the case with VAE training in general. It all boils down to the task you are working on. If you want your model to generate meaningful reconstructions and it does that then I wouldn't worry. If you want to generate novel samples and that is failing maybe some hyper-parameter tuning is needed.

Specifically, for your case you might also want to try early-stopping. As I understand from the log-file KL starts increasing and thus affecting your overall loss after epoch 220. It also seems to me that your reconstruction loss is stable around that epoch mark.

In terms of battling KL divergence issues I would also refer you to the beta-VAE due to Higgins et. al (https://openreview.net/forum?id=Sy2fzU9gl). In a nutshell they add a term (beta) that regularizes the KL term of the ELBO. You can experiment with that and see how it affects training.

Finally, there is a very helpful post from MSR where they introduce annealing regimes for the the beta hyper-parameter in beta-VAEs and how they affect training (https://www.microsoft.com/en-us/research/blog/less-pain-more-gain-a-simple-method-for-vae-training-with-less-of-that-kl-vanishing-agony/ ).

beta annealing

$\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.