Jigsaw-VAE: Towards Balancing Features in Variational Autoencoders
This work addresses a specific problem in unsupervised learning for researchers and practitioners using VAEs, offering an incremental improvement to enhance feature generalization and sample diversity.
The paper tackles the problem of feature imbalance in Variational Autoencoders (VAEs), where latent variables focus on some factors of variation at the expense of others, leading to poor generalization and less diverse generated samples. They propose a regularization scheme that substantially addresses this issue and introduce a metric to measure feature balance in generated images.
The latent variables learned by VAEs have seen considerable interest as an unsupervised way of extracting features, which can then be used for downstream tasks. There is a growing interest in the question of whether features learned on one environment will generalize across different environments. We demonstrate here that VAE latent variables often focus on some factors of variation at the expense of others - in this case we refer to the features as ``imbalanced''. Feature imbalance leads to poor generalization when the latent variables are used in an environment where the presence of features changes. Similarly, latent variables trained with imbalanced features induce the VAE to generate less diverse (i.e. biased towards dominant features) samples. To address this, we propose a regularization scheme for VAEs, which we show substantially addresses the feature imbalance problem. We also introduce a simple metric to measure the balance of features in generated images.