1
$\begingroup$

I'm applying VAEs to sections genomic data (haplotypic vcf format, so binary variables), with one model being trained on each section. They each have different layer sizes and weights to better fit the sections of data, but share training hyperparametres. Depending on which section I'm applying these methods to, my models sometimes exhibit mode searching behaviour quite heavily.

PCA plots showing data before and after autoencoding

These are PCA plots of my data for 4 sections before and after autoencoding. The base values are represented in grey, and the decoded values in red. So a perfect autoencoding process should have one red dot over each grey dot (I know that isn't the true goal of a VAE).

In the two plots on the left, the decoded data quite visibly concentrates in the densest regions of the distribution of the base data, with a "path" of samples between these modes. On the right, the decoded data is much more dispersed throughout the distribution of the base data. To obtain this, I already greatly decreased the weight of the KL divergence within the loss calculations of my VAE models to allow them to "explore" the latent space more, but this seems to not be sufficient to eliminate the mode searching behaviour in all of them.

What paths can I explore to try to remedy this mode searching behavious in my VAE models? I am working with keras in Python3, so bonus points if your answer includes methods that work in that environment.

$\endgroup$

1 Answer 1

1
$\begingroup$

VAE objective is to maximise the ELBO: $$ \int_x p(x) \int_z q(z|x) \big\{ \log p(x|z) - \log \tfrac{q(z|x)}{p(z)} \big\} $$ the first term reconstructs the data, the second term maximises the entropy of $q(z|x)$ ("puffing out" the posteriors) and constrains latent representations $z$ to fit the prior $p(z)$. There are few parameters to tweak,

  • the prior $p(z)$ is typically chosen as a standard Gaussian,
  • $q(z|x)$ learns to approximate the true posterior $p(z|x) \doteq \tfrac{p(x|z)p(z)}{\int_z p(x|z)p(z)}$
  • reconstruction is driven by maximising $p(x|z) =\mathcal{N}(x; dec(z), \sigma_x^2)$

The main parameter to consider is the reconstruction variance $\sigma_x^2$. Setting this to be smaller lowers the probability of poor reconstructions, which must improve to maximise the term (if the network architecture is flexible enough to allow it).

Re generation: while reconstructions should improve (for training set), if the networks are not sufficiently flexible, the distribution of $z$ may not fit $p(z)$ (e.g. consider if both encoder and decoder were linear). So sampling $p(z)$ may give $z$'s that don't correspond to the data and generating from them gives gibberish.

[Note: $\sigma_x^2$ is in the denominator of the (log) Gaussian reconstruction term, so optimising the ELBO by SGD is unchanged by multiplying through by $\sigma_x^2$ (with an adjusted learning rate), so it is common to use an MSE reconstruction term (w/o $\sigma_x^2$) and "weight" the KL term instead. This weight can be interpreted as $\sigma_x^2$ and reducing it has the same effect as above.]

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.