Affine Variational Autoencoders

Variational autoencoders (VAEs) have in recent years become one of the most powerful approaches to learning useful latent representations of data in an unsupervised manner. However, a major challenge with VAEs is that they have tremendous difficulty in generalizing to data that deviate from the training set (e.g., perturbed image variants). Normally data augmentation is leveraged to overcome this limitation; however, this is not only computational expensive but also necessitates the construction of more complex models. In this study, we introduce the notion of affine variational autoencoders (AVAEs), which extends upon the conventional VAE architecture through the introduction of affine layers. More specifically, within the AVAE architecture an affine layer perturbs the input image prior to the encoder, and a second affine layer performs an inverse perturbation to the output of the decoder. The parameters of the affine layers are learned to enable the AVAE to encode images at canonical perturbations, resulting in a better reconstruction and a disentangled latent space without the need for data augmentation or the use of more complex models. Experimental results demonstrate the efficacy of the proposed VAE architecture for generalizing to images in the MNIST validation set under affine perturbations without the need for data augmentation, demonstrating significantly reduced loss when compared to conventional VAEs.