# how to weight KLD loss vs reconstruction loss in variational auto-encoder

in nearly all code examples I've seen of a VAE, the loss functions are defined as follows (this is tensorflow code, but I've seen similar for theano, torch etc. It's also for a convnet, but that's also not too relevant, just affects the axes the sums are taken over):

# latent space loss. KL divergence between latent space distribution and unit gaussian, for each batch.
# first half of eq 10. in https://arxiv.org/abs/1312.6114
kl_loss = -0.5 * tf.reduce_sum(1 + log_sigma_sq - tf.square(mu) - tf.exp(log_sigma_sq), axis=1)

# reconstruction error, using pixel-wise L2 loss, for each batch
rec_loss = tf.reduce_sum(tf.squared_difference(y, x), axis=[1,2,3])

# or binary cross entropy (assuming 0...1 values)
y = tf.clip_by_value(y, 1e-8, 1-1e-8) # prevent nan on log(0)
rec_loss = -tf.reduce_sum(x * tf.log(y) + (1-x) * tf.log(1-y), axis=[1,2,3])

# sum the two and average over batches
loss = tf.reduce_mean(kl_loss + rec_loss)


However the numeric range of kl_loss and rec_loss are very dependent on latent space dims and input feature size (e.g. pixel resolution) respectively. Would it be sensible to replace the reduce_sum's with reduce_mean to get per z-dim KLD and per pixel (or feature) LSE or BCE? More importantly, how do we weight latent loss with reconstruction loss when summing together for the final loss? Is it just trial and error? or is there some theory (or at least rule of thumb) for it? I couldn't find any info on this anywhere (including the original paper).

The issue I'm having, is that if the balance between my input feature (x) dimensions and latent space (z) dimensions is not 'optimum', either my reconstructions are very good but the learnt latent space is unstructured (if x dimensions is very high and reconstruction error dominates over KLD), or vice versa (reconstructions are not good but learnt latent space is well structured if KLD dominates).

I'm finding myself having to normalise reconstruction loss (dividing by input feature size), and KLD (dividing by z dimensions) and then manually weighting the KLD term with an arbitrary weight factor (The normalisation is so that I can use the same or similar weight independent of dimensions of x or z). Empirically I've found around 0.1 to provide a good balance between reconstruction and structured latent space which feels like a 'sweet spot' to me. I'm looking for prior work in this area.

Upon request, maths notation of above (focusing on L2 loss for reconstruction error)

$$\mathcal{L}_{latent}^{(i)} = -\frac{1}{2} \sum_{j=1}^{J}(1+\log (\sigma_j^{(i)})^2 - (\mu_j^{(i)})^2 - (\sigma_j^{(i)})^2)$$

$$\mathcal{L}_{recon}^{(i)} = -\sum_{k=1}^{K}(y_k^{(i)}-x_k^{(i)})^2$$

$$\mathcal{L}^{(m)} = \frac{1}{M}\sum_{i=1}^{M}(\mathcal{L}_{latent}^{(i)} + \mathcal{L}_{recon}^{(i)})$$

where $J$ is the dimensionality of latent vector $z$ (and corresponding mean $\mu$ and variance $\sigma^2$), $K$ is the dimensionality of the input features, $M$ is the mini-batch size, the superscript $(i)$ denotes the $i$th data point and $\mathcal{L}^{(m)}$ is the loss for the $m$th mini-batch.

Namely:

beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework

discusses my exact question with a few experiments. Interestingly, it seems their $\beta_{norm}$ (which is similar to my normalised KLD weight) is also centred around 0.1, with higher values giving more structured latent space at the cost of poorer reconstruction, and lower values giving better reconstruction with less structured latent space (though their focus is specifically on learning disentangled representations).

and related reading (where similar issues are discussed)

I would like to add one more paper relating to this issue (I cannot comment due to my low reputation at the moment).

In subsection 3.1 of the paper, the authors specified that they failed to train a straight implementation of VAE that equally weighted the likelihood and the KL divergence. In their case, the KL loss was undesirably reduced to zero, although it was expected to have a small value. To overcome this, they proposed to use "KL cost annealing", which slowly increased the weight factor of the KL divergence term (blue curve) from 0 to 1. This work-around solution is also applied in Ladder VAE.

Paper:

Bowman, S.R., Vilnis, L., Vinyals, O., Dai, A.M., Jozefowicz, R. and Bengio, S., 2015. Generating sentences from a continuous space. arXiv preprint arXiv:1511.06349.