Replacing the KL-divergence term in a VAE with parameter regularization

Question

When training a VAE, one aim to optimize function $\mathcal{L}$, defined as:

$$\mathcal{L}\left(\theta,\phi; \mathbf{x}^{(i)}\right) = - D_{KL}\left(q_\phi(\mathbf{z}|\mathbf{x}^{(i)}) || p_\theta(\mathbf{z})\right) + \mathbb{E}_{q_\phi\left(\mathbf{z}|\mathbf{x}^{(i)}\right)}{\log p_\theta\left(\mathbf{x}^{(i)}| \mathbf{z}\right)},$$ where $q_\phi\left(\mathbf{z}|\mathbf{x}^{(i)}\right) = \mathcal{N}_J(\mathbf{\mu}^{(i)},\sigma^{(i)}\mathbb{1})$ and $p_\theta(\mathbf{z}) = \mathcal{N}_J(\mathbf{0},\mathbb{1})$. The term $D_{KL}\left(q_\phi(\mathbf{z}|\mathbf{x}^{(i)}) || p_\theta(\mathbf{z})\right)$ may be viewed as a regularizer.

Since normality is already imposed on $q_\phi\left(\mathbf{z}|\mathbf{x}^{(i)}\right)$, then the KL-divergence term aims solely to turn $\mathbf{\mu}^{(i)}=0$ and $\mathbf{\sigma}^{(i)}=1$. Thus, wouldn't it be equivalent (and simpler) to replace the $D_{KL}$ term by (something along the lines of) $\left(\mu^{(i)}\right)^2 + \left(\sigma^{(i)}-1\right)^2$? That is, to use the following function $\tilde{\mathcal{L}}$ instead:

$$\tilde{\mathcal{L}}\left(\theta,\phi; \mathbf{x}^{(i)}\right) = -\sum_{j=1}^J\left(\left(\mu^{(i)}_j\right)^2 + \left(\sigma^{(i)}_j-1\right)^2\right) + \mathbb{E}_{q_\phi\left(\mathbf{z}|\mathbf{x}^{(i)}\right)}{\log p_\theta\left(\mathbf{x}^{(i)}| \mathbf{z}\right)}.$$

Answer 1

Doing what you propose, i.e. adding a regularization term other than the KLD in the loss, is totally feasible. You can find many classical autoencoder architectures which incorporate a diversity of regularization term along with the regularization term, see the wiki entry for autoencoder for example.

Now the critical point is that you will not be correct referring to your new model as a VAE because a VAE is strictly derived in the variational Bayes setting by constructing a lower bound on the loglikelihood. If you replace the regularizing KLD term without proper care, you bindly modify the loss and cannot assert that you still have an ELBO whose maximization is at the core of VAE training.

Answer 2

The VAE is a direct implementation of the ELBO, a lower bound on the log likelihood of the data, which makes optimising its training objective a provably valid approach. Any deviation from that would need to be considered carefully as you may otherwise have no idea what optimising the new objective would achieve, i.e. what the model would learn.

• Aim: train a model of the data $p_\theta(x)$ to fit the true data distribution $p(x)$, where the model is defined in terms of latent variable $z$: $p_\theta(x) = \int p_\theta(x|z)p_\theta(z) dz$
• Approach: minimise the KL divergence between $p_\theta(x)$ and $p(x)$
  • i.e. find $\text{argmin}_{\theta} \int_x p(x)\log\tfrac{p(x)}{p_\theta(x)} = \text{argmax}_\theta \int_x p(x)\log p_\theta(x)$
  • i.e. maximising cross entropy (RHS) is equivalent to minimising the KL divergence (LHS)
  • it is typically intractable to take gradients w.r.t $\theta$ to maximise this, so a lower bound is maximised:
    • $\int p(x)\log p_\theta(x) \geq \int_x p(x)\int_z q(z|x)\{\log p_\theta(x|z) - \tfrac{q(z|x)}{p_\theta(z)}\}\quad$ [ = ELBO in "VAE form"]
• While the terms of this expression can be thought of intuitively as reconstruction loss + regulariser, this is just intuition, the fact is you want $p_\theta(x)$ to learn $p(x)$, which maximising the ELBO guarantees. If you "hack around with it", that guarantee may be lost.
• The proposed alternative "regulariser" may (i) happen to equate to $\tfrac{q(z|x)}{p_\theta(z)}$ for particular choices of those distributions, in which case you are baking in those assumptions (which may be good or bad depending on the data distribution); or (ii) to some extent approximate particular distributional assumptions and "seem to work" (e.g. on a given test set), but, given it is an approximation (and so "wrong" in places) there may be regions in the data space where the VAE performs poorly.