I am training VAE autoencoder model. VAE has loss combining MSE+KL divergence. When I train the model, KL loss is increasing over or near 100 while MSE loss is decreasing. So, can anyone tell me what cause for that and how to reduce the KL loss? My data is coordinate data.
2 Answers
This is standard behavior for VAEs. The thing that ultimately matters is that the total loss (reconstruction error plus KL divergence) is decreasing. If the model is good enough to solve your problem, it's a success.
Sometimes a VAE will have the KL divergence swamp any improvement to the reconstruction. In that case, it can help to anneal the weight assigned to the KLD portion of the loss from 0 (KLD is ignored) to 1 (KLD is given full weight and the loss is the ordinary variational lower-bound).
Samuel R. Bowman, Luke Vilnis. "Generating Sentences from a Continuous Space"
I am using VAE for the dimensionality reduction of 1D data. I have also experienced, the sudden rise of KLD, which I thought would decrease anyway. See the image below:
I am using Gaussian likelihood as a reconstruction error instead and it also increases. But
$$ \beta * KL - reconstruction $$
as a total loss function decreases which @Sycorax mentioned. See the total loss below:
The $\beta$ comes from a paper called 'disentangled VAE', higher value reduces the performance of the decoder generating poor reconstructions. So a balance should be searched based on your application/data.