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 KL-Divergence = 0 in all epochs.

KL-Divergence normal KL-Divergence and ReconstrunctionLoss with same weight.

KL-Divergence weight dynamic 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! :)

  • $\begingroup$ 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. $\endgroup$
    – mesolmaz
    Jul 1 at 18:32

Update: I already found the error. It was pretty silly actually. The point is that, when training VAEs, you should compute the sum error per each feature, and not the average as it's usual. This means to use the reduction='sum' parameter.

torch.nn.MSELoss(reduction='sum') https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html

If this is not done, the reconstruction error will stuck up in a point of the search space in which the average MSE Loss is good for all the features, but not good enough to accurately reconstruct the data.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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