Jigsaw-VAE: Towards Balancing Features in Variational Autoencoders

TL;DR

This paper introduces Jigsaw-VAE, a regularization method for variational autoencoders that balances learned features, improving their generalization across environments and diversity of generated samples.

Contribution

The paper proposes a novel regularization scheme for VAEs to address feature imbalance, enhancing feature diversity and cross-environment generalization.

Findings

Regularization reduces feature imbalance in VAEs.

Balanced features improve generalization across environments.

Enhanced diversity in generated samples.

Abstract

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.