# 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)

• Did you end up using reduce_mean instead of reduce_sum? I find it that when I use reduce_sum, reconstruction loss is 2-3 orders of magnitude larger than KL loss. But when I use reduce_mean (they are the same order of magnitude), reconstruction is basically just black images. (Although I haven’t let the training run for too long.) Dec 9, 2020 at 10:33

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.

I faced the same problem of not knowing how to weigh the reconstruction and KL terms, and here I want to add an answer with some concrete values of $$\beta$$ (the weight of KL term) for future reference.

In the $$\beta$$-VAE paper, they seem to use the values of $$\beta_{norm}$$ ranging between $$0.001$$ and $$10$$ (Fig. 6 from the paper).

They calculate $$\beta_{norm}$$ as follows:

$$\beta_{norm} = \frac{\beta M}{N},$$

where $$M$$ is the size of latent space (e.g. $$10$$) and $$N$$ is the input size (e.g. $$64 \cdot 64 \cdot 1 = 4096$$).

Update on Dec. 6th 2020: I made a blog post to explain this in details.

I finally manage to figure out the reason of weighting KL divergence in VAE. It is more about the normalized constant of the distribution modeled the target variable. Here, I am going to present some output distributions we often use. Most of the notation will follow the book "Pattern recognitions and Machine learning".

1. Linear regression (unbounded regression): (section 3.1.1 on page 140) - This explains for the weighting KL divergence when using MSE loss

The target variable $$t$$ is assumed to be the sum of the deterministic function $$y(\mathbf{x}, \mathbf{w})$$ and a Gaussian noise: $$\begin{equation} t = y(\mathbf{x}, \mathbf{w}) + \epsilon, \qquad\epsilon \sim \mathcal{N}\left(\epsilon | 0, \color{red}{\beta}^{-1}\right) \end{equation}$$ The target variable is therefore modeled as a Gaussian random variable with the log-likelihood given as: $$\begin{equation} p(t | \mathbf{x}, \mathbf{w}, \color{red}{\beta}) = \mathcal{N} \left( t | y(\mathbf{x}, \mathbf{w}), \color{red}{\beta}^{-1} \right) \end{equation}$$ Given this assumption, the log-likelihood at data points $$\{\mathbf{x}_{n}, t_{n}\}_{n=1}^{N}$$ is: $$\begin{equation} \ln p(\mathbf{t} | \mathbf{x}, \mathbf{w}, \color{red}{\beta}) = \frac{N}{2} \ln \frac{\color{red}{\beta}}{2\pi} - \color{red}{\beta} E_{D}(\mathbf{w}), \end{equation}$$ where: $$\begin{equation} E_{D}(\mathbf{w}) = \frac{1}{2} \sum_{n=1}^{N} [t_{n} - y(\mathbf{x}, \mathbf{w})]^{2}. \end{equation}$$ We often optimize only $$E_{D}(\mathbf{w})$$, not the whole log-likelihood $$\ln p(\mathbf{t} | \mathbf{x}, \mathbf{w}, \beta)$$, resulting in ignoring the precision $$\color{red}{\beta}$$. This might be fine for conventional regression where the loss consists of only the negative log-likelihood (NLL) $$-\ln p(\mathbf{t} | \mathbf{x}, \mathbf{w}, \beta)$$, and the prediction would be the mean of the target variable $$t$$.

However, the loss in VAE consists of the NLL (or reconstruction loss) and the regularization (KL loss). Therefore, if the weight factor of MSE term (or, $$E_{D}(\mathbf{w})$$ in this case) is 1, we need to weight the KL divergence with a factor $$\beta_{KL} = 1/\color{red}{\beta}$$ to be mathematically correct. In practice, people often find a good value of the precision $$\beta_{KL}$$ through hyper-parameter tuning. Another approach is to learn $$\color{red}{\beta}$$ by considering it as a learnable parameter which is obtained by minimizing the whole VAE loss function.

1. Logistic regression - This explains the case of binary cross-entropy loss used for black-and-white images

Let's consider the case of binary classification. The ground-truth is either 0 or 1, and the target variable $$t = p(y = 1 | \mathbf{x})$$ is assumed to follow a Bernoulli distribution: $$\begin{equation} p(t | \mathbf{x}, \mathbf{w}) = \mathcal{B}(t | y(\mathbf{x}, \mathbf{w})) = \left[y(\mathbf{x}, \mathbf{w})\right]^{t} \left[ 1 - y(\mathbf{x}, \mathbf{w}) \right)^{1 - t}. \end{equation}$$ Hence, the NLL in this case is given by: $$\begin{equation} -\ln p(t | \mathbf{x}, \mathbf{w}) = -\left[ t \ln y(\mathbf{x}, \mathbf{w}) + (1 - t) \ln (1 - y(\mathbf{x}, \mathbf{w})) \right], \end{equation}$$ which is the binary cross-entropy loss. (One can extend to softmax for multiclass classification by using a categorical distribution to lead to cross-entropy loss.)

For MNIST (or black and white images) data set, each pixel is either 0 or 1, and therefore, we can use binary cross-entropy loss as the reconstruction loss in the VAE to predict the probability that the value of a pixel is 1. And since the mean of the Bernoulli distribution equals to $$y(\mathbf{x}, \mathbf{w})$$, we often use $$y(\mathbf{x}, \mathbf{w})$$ as pixel intensity to plot the reconstructed images.

Note that when using binary cross-entropy loss in a VAE for black and white images, we do not need to weight the KL divergence term, which has been seen in many implementations.

1. Bounded regression (e.g. regression in [0, 1]) - This explains the case of weighting KL divergence when using binary cross-entropy loss for color images

As explained in logistic regression, the support (or the label) of a Bernoulli distribution is $$\{0, 1\}$$, not $$[0, 1]$$, but in practice, it is still employed for color-image reconstruction, which requires a support in $$[0, 1]$$, or $$\{0, 1, \ldots, 255\}$$.

Since our interest is the case for support in $$[0, 1]$$, we could find some continuous distribution that has support in $$[0, 1]$$ to model our prediction. One simple one is the beta distribution. In that case, our prediction would be the 2 parameters $$\alpha$$ and $$\beta$$. Seem complicated?

Fortunately, a continuous version of Bernoulli distribution has been proposed recently, so that, we can still use the binary cross-entropy loss to predict the intensive of a pixel with some minimal modification. Please refer to the paper "The continuous Bernoulli distribution: fixing a pervasive error in VAE" or Wikipedia page for further details of the distribution.

Under the assumption of the continuous Bernoulli distribution, the log-likelihood can be expressed as: $$\begin{equation} \ln p(t | \mathbf{x}, \mathbf{w}) = \mathcal{CB}(t | y(\mathbf{x}, \mathbf{w})) = C(y(\mathbf{x}, \mathbf{w})) (y(\mathbf{x}, \mathbf{w}))^{t} (1 - y(\mathbf{x}, \mathbf{w}))^{1-t}, \end{equation}$$ where $$C(y(\mathbf{x}, \mathbf{w}))$$ is the normalized constant.

Hence, when working with VAE involving binary cross-entropy loss, instead of tuning for a weight factor of the KL term - which might be mathematically incorrect, we simply add $$- \ln C(y(\mathbf{x}, \mathbf{w}))$$ into the loss, and then optimize.

I want to add another interesting paper relating to this question, where the authors propose a cyclical annealing scheme for the KLD term to improve the training of a VAE for natural language processing tasks.

The basic idea is to repeat the process of increasing the weighting term $$\beta$$ multiple times starting from zero. The authors argue that this warm re-start allows one to learn more meaningful latent codes progressively by leveraging the results of previous learning cycles. Furthermore, they show that this strategy can improve reconstruction loss compared to a monotonically increasing $$\beta$$ only at the beginning of the training. $\beta$ schedules" />

Paper:

Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin. 2019. Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 240–250, Minneapolis, Minnesota. Association for Computational Linguistics. DOI: 10.18653/v1/N19-1021

Code:

github.com/haofuml/cyclical_annealing