AAVAE: Augmentation-Augmented Variational Autoencoders

William Alejandro Falcon, Ananya Harsh Jha, Teddy Koker, Kyunghyun Cho

Keywords: Self-Supervised Learning, Autoencoders, Variational Autoencoders, Data Augmentation

Abstract: Recent methods for self-supervised learning can be grouped into two paradigms: contrastive and non-contrastive approaches. Their success can largely be attributed to data augmentation pipelines which generate multiple views of a single input that preserve the underlying semantics. In this work, we introduce augmentation-augmented variational autoencoders (AAVAE), yet another alternative to self-supervised learning, based on autoencoding. We derive AAVAE starting from the conventional variational autoencoder (VAE), by replacing the KL divergence regularization, which is agnostic to the input domain, with data augmentations that explicitly encourage the internal representations to encode domain-specific invariances and equivariances. We empirically evaluate the proposed AAVAE on image classification, similar to how recent contrastive and non-contrastive learning algorithms have been evaluated. Our experiments confirm the effectiveness of data augmentation as a replacement for KL divergence regularization. The AAVAE outperforms the VAE by 30% on CIFAR-10, 40% on STL-10 and 45% on Imagenet. On CIFAR-10 and STL-10, the results for AAVAE are largely comparable to the state-of-the-art algorithms for self-supervised learning.

One-sentence Summary: We show the inadequacy of KL-divergence as a regularizer for learning representations with VAEs. Instead, autoencoders with a data augmentation pipeline, sampling in the latent space, and the absence of a KL-divergence term prove to be better.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2107.12329/code)