2
$\begingroup$

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! :)

Improve this question
$\endgroup$
1
  • 1
    $\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
    Commented Jul 1, 2021 at 18:32

1 Answer 1

Reset to default
2
$\begingroup$

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.

Improve this answer
$\endgroup$
2
  • $\begingroup$ I really don't understand your comment here @Ale $\endgroup$
    – Vladimir Vargas
    Commented Mar 10, 2022 at 0:32
  • $\begingroup$ 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? $\endgroup$
    – lucidbrot
    Commented Jun 2, 2023 at 11:18

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.