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.
$\beta$ schedules" />