VAE Loss: The weight of BCE vs. KL

[huanghoujing]

The BCE loss is averaged over the batch dimension, but KL is averaged over both batch and image pixel dimension. I hope to know the reason for this.

I tried both

• averaging the BCE loss over both dimensions
• averaging the KL loss just over batch dimension

But the reconstructed images and generated images in test time are rather faint and blurry (just like digit 8), which is far worse than the setting here.

If the setting here is just to balance between two loss, why do you use the number of image pixels (which is intuitively supposed to have some mathematical meaning) for scaling. Why not use some number like 10, 100 or 1000?

Thanks for help.

[helperfunc]

According to torch.nn.BCELoss doc, and its cpu implementation source code , if size_average=True, BCE loss is averaged over both batch and image pixel dimension. Because when computing - log(input) * target - log(1 - input) * (1 - target), both input and target are resized into one dimension via view(-1), and the outcome is divided by the one dimension size, which is batch_size * image pixel dimension.

[huanghoujing]

I am still confused about how we should set the weight between two losses.

Anyway, thank you very much!