Data Science Stack Exchange is a question and answer site for Data science professionals, Machine Learning specialists, and those interested in learning more about the field. It only takes a minute to sign up.
Sign up to join this community
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.
I have attached link to training logs file.
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/ ).
