mse loss is way much higher than the kl loss in eq (3), how do you deal with the imbalance?

[kailigo]

Hi I am trying to reimplement the code. My current implementation shows that
in eq (3), the mse loss is way much higher than the kl loss in eq (3) (with 4 magnitude difference)
And you use the same weight for the two losses, which results in nan values for other losses, as the mse loss dominant the overall learning objective.
Thanks for the answer in advance.

[lryta]

The camera-ready version has changed both losses to kld losses. MSE loss is basically an implementation error that was found too late which got into the first submitted version.