Constraining Variational Inference with Geometric Jensen-Shannon Divergence

TL;DR

This paper introduces a novel regularization method for variational autoencoders using a skew-geometric Jensen-Shannon divergence, improving reconstruction and generation quality by better controlling divergence in latent space.

Contribution

It proposes a new divergence-based regularization mechanism for VAEs that interpolates between forward and reverse KL, enhancing latent space control and generative performance.

Findings

Improved reconstruction quality over baseline VAEs.

Enhanced generative capabilities demonstrated through experiments.

Simple hyperparameter with clear latent space interpretation.

Abstract

We examine the problem of controlling divergences for latent space regularisation in variational autoencoders. Specifically, when aiming to reconstruct example $x\in\mathbb{R}^{m}$ via latent space $z\in\mathbb{R}^{n}$ ($n\leq m$), while balancing this against the need for generalisable latent representations. We present a regularisation mechanism based on the skew-geometric Jensen-Shannon divergence $\left(\textrm{JS}^{\textrm{G}_{\alpha}}\right)$. We find a variation in $\textrm{JS}^{\textrm{G}_{\alpha}}$, motivated by limiting cases, which leads to an intuitive interpolation between forward and reverse KL in the space of both distributions and divergences. We motivate its potential benefits for VAEs through low-dimensional examples, before presenting quantitative and qualitative results. Our experiments demonstrate that skewing our variant of $\textrm{JS}^{\textrm{G}_{\alpha}}$, in the context of $\textrm{JS}^{\textrm{G}_{\alpha}}$-VAEs, leads to better reconstruction and generation when compared to several baseline VAEs. Our approach is entirely unsupervised and utilises only one hyperparameter which can be easily interpreted in latent space.