# Variational Autoencoder: balance KL-Divergence and ReconstructionLoss

I'm facing some problems working with VAE. They are pretty similar than the ones at this thread: Balancing Reconstruction vs KL Loss Variational Autoencoder but I still have the problems and I don't really understand them, that's why I open this new thread.

The point is that it looks like the Reconstruction Loss and the KL-Divergence are opposite terms. I've tried many different things, and here are the graphs of the loss:

KL-Divergence = 0 in all epochs.

KL-Divergence and ReconstrunctionLoss with same weight.

The weight of the KL-Divergence change from 0 to 1 progressively.

As you can see in the graphs, if the KL_Divergence is = 0, the ReconstructionLoss improves. But if they have the same weight, the ReconstrunctionLoss is always the same and it only improves the KL-Divergence. Actually, while KL-Divergence = 0 and the ReconstructionLoss improves, the real KL-Divergence worsen.

Any idea about how to solve this? I follow what appears in this paper about KL cost annealing: https://arxiv.org/pdf/1511.06349.pdf, and the result is my third chart.

I'm completely lost and have no idea about how to continue working with this.

One last thing: I'm not working with images, but to visualize my data I use UMAP. The UMAP plot of the original data is completely different from the UMAP plot of the reconstructed data if I don't set KL-Divergence equal to 0. Using KL-Divergence as part of loss function is just worsen the perfomance of my VAE...

Thanks you so much! :)

• Do you mind sharing your code? Sometimes reconstruction loss dominates the KL loss but you should expect reconstruction loss to decrease while KL loss increases as the training progresses. Commented Jul 1, 2021 at 18:32

• I don't understand, why would 'sum' vs 'mean' have a relevant difference except in the scale of the loss value and that the sum version might vary based on batch size or input size? Commented Jun 2, 2023 at 11:18