AAVAE: Augmentation-Augmented Variational Autoencoders

1 INTRODUCTION . The goal of self-supervised learning ( SSL ) ( DeSa , 1993 ) is to learn good representations from unlabeled examples . A good representation is often defined as the one that reflects underlying class structures well . The quality of a representation obtained from SSL is evaluated by measuring downstream classification accuracy on a labelled dataset . In recent years , two families of approaches have emerged as the state-of-the-art for SSL : contrastive and non-contrastive learning . At its core , a contrastive learning algorithm stochastically creates two views from each training example , called positive and anchor examples , and selects one of the other training examples as a negative ( Bachman et al. , 2019 ; Tian et al. , 2019 ; Chen et al. , 2020b ; a ) . The positive and anchor examples are brought closer in the representation space , while the negative example is pushed away from the anchor . This definition of contrastive loss brings in two interconnected issues . First , there is no principled way to choose negative examples , and hence these negatives are chosen somewhat arbitrarily each time . Second , the contrastive loss function is not decomposed over training examples because negative examples come from other training examples . Partly to address these limitations , recent studies have proposed non-contrastive approaches that have removed the need for negative examples ( Grill et al. , 2020 ; Caron et al. , 2020 ; Zbontar et al. , 2021 ) . These approaches avoid the necessity of explicit negatives by constraining or regularizing dataset-level statistics of internal representation ( Zbontar et al. , 2021 ; Caron et al. , 2021 ; Genevay et al. , 2018 ) . Dataset-level statistics , which are intractable to compute , are instead approximated using a minibatch of training examples . This often results in the need of large minibatches . Also , the use of batch-level statistics means that non-contrastive losses are not decomposable as well . Despite the apparent differences between these two families of algorithms , they all recognize the importance of and rely heavily on data augmentation as a way of incorporating domain knowledge . For instance , Chen et al . ( 2020a ) have highlighted that the downstream accuracy after finetuning varied between 2.6 % and 69.3 % on ImageNet ( Deng et al. , 2009 ) , depending on the choice of data augmentation . This is perhaps unsurprising since the importance of domain knowledge has been reported in various domains beyond computer vision . In reinforcement learning , Kostrikov et al . ( 2020 ) and Raileanu et al . ( 2020 ) have shown the benefit of adding domain information via pixel-level data augmentation in continuous control . In natural language processing , Ng et al . ( 2020 ) demonstrate the effectiveness of domain-specific augmentation by using a pretrained denoising autoencoder to build a robust classifier . A variational autoencoder ( VAE ) implements a latent variable model using a composition of two neural networks . A neural net decoder maps a latent variable configuration to an observation , and a neural net encoder approximately infers the latent variable configuration given the observation ( Kingma & Welling , 2013 ) . It is often trained to maximize the variational lowerbound or its variant ( Kingma & Welling , 2013 ; Higgins et al. , 2016 ) . Careful inspection of this learning objective shows two parts : autoencoding and latent-space regularization . Autoencoding ensures that there is an approximately one-to-one mapping between individual inputs and internal representations . This prevents the collapse of internal representations onto a single point , similar to what negative examples in contrastive learning and regularization of batch-level statistics in non-contrastive learning do . Latent-space regularization , on the other hand , ensures that the internal representation is arranged semantically in a compact subset of the space . It is often done by minimizing the KL divergence ( Kullback & Leibler , 1951 ) from the approximate posterior , returned by the encoder , to the prior distribution and adding noise to the representation during training ( i.e. , sampling from the approximate posterior ) . This performs a role similar to that of data augmentation in contrastive and non-contrastive approaches but is different in a way that it is agnostic to the input domain . Based on these observations : ( 1 ) the importance of data augmentations and ( 2 ) variational autoencoders for representation learning , we propose a third family of self-supervised learning algorithms in which we augment variational autoencoders with data augmentation . We refer to this family of models as Augmentation-Augmented Variational Autoencoders ( AAVAE ) . In AAVAEs , we replace the usual KL-divergence ( Kullback & Leibler , 1951 ) term in ELBO ( Kingma & Welling , 2013 ) with a denoising criterion ( Vincent et al. , 2008 ; 2010 ) based on domain-specific data augmentation . We hypothesize that this new approach allows the representations learned by AAVAEs to encode domain-specific data invariances and equivariances . The resulting model offers a few advantages over the existing contrastive and non-contrastive methods . First , the loss function is not dependent on the batch-level statistics , which we suspect enables us to use smaller minibatches . Second , the AAVAE does not necessitate an arbitrary choice of negative sampling strategy . We pretrain AAVAEs on image datasets : CIFAR-10 ( Krizhevsky & Hinton , 2009 ) , STL-10 ( Coates et al. , 2011 ) and Imagenet ( Deng et al. , 2009 ) , and as is the norm with other recently proposed approaches ( Goyal et al. , 2019 ; Chen et al. , 2020a ; Caron et al. , 2020 ) , we evaluate them on classification tasks corresponding to the dataset using a single linear layer without propagating gradients back to the encoder . We find that our autoencoding-based method gives a downstream classification performance comparable to the current state-of-the-art SSL methods , with 87.14 % accuracy on CIFAR-10 and 84.72 % on STL-10 . On Imagenet , the AAVAE outperforms the carefully crafted pretext tasks for SSL , such as Colorization ( Zhang et al. , 2016 ) , Jigsaw ( Noroozi & Favaro , 2016 ) and Rotation ( Gidaris et al. , 2018 ) , demonstrating that designing such complex pretext tasks is unnecessary . As anticipated from our formulation , representation learned by the AAVAE is robust to the choice of hyperparameters , including minibatch size , latent space dimension , and the network architecture of the decoder . Our observations strongly suggest that autoencoding is a viable third family of self-supervised learning approach in addition to contrastive and non-contrastive learning . 2 SELF-SUPERVISED LEARNING . Self-supervised learning ( SSL ) aims to derive training signal from the implicit structure present within data ( DeSa , 1993 ) . This enables SSL methods to leverage large unlabeled datasets to learn representations ( Goyal et al. , 2021 ) which can then be used to solve downstream tasks , such as classification and segmentation , for which it is often expensive to collect a large number of annotations . Here , we summarize quite a few variations of this approach proposed over the last few years . 2.1 PRETEXT TASKS . Pretext tasks are designed to train a neural network to predict a non-trivial but easily applicable transformation applied to the input . For example , Gidaris et al . ( 2018 ) randomly rotate an input image by 0◦ , 90◦ , 180◦ , or 270◦ and train a network to predict the angle of rotation . The colorization pretext task ( Zhang et al. , 2016 ) creates a training signal by converting RGB images to grayscale and training a network to restore the removed color channels . Image inpainting ( Pathak et al. , 2016 ) learns representations by training an encoder-decoder network to fill in artificially-occluded parts of an image . Both jigsaw ( Noroozi & Favaro , 2016 ) and relative patch prediction ( Doersch et al. , 2015 ) tasks divide an input image into patches . The jigsaw task ( Noroozi & Favaro , 2016 ) shuffles the spatial ordering of these patches and trains a network to predict the correct order . In contrast , relative patch prediction ( Doersch et al. , 2015 ) selects two patches of an image and asks the network to predict their relative spatial positions . More recently , Doersch & Zisserman ( 2017 ) combined various pretext tasks into a single method . Goyal et al . ( 2019 ) have , however , shown that training neural network backbones using pretext tasks often does not capture representations invariant to pixel-space perturbations . Consequently , these representations perform poorly on downstream tasks while they solve the original pretext task well . 2.2 CONTRASTIVE LEARNING . Between the two major families of state-of-the-art methods for self-supervised learning , we discuss the one based on the so-called contrastive loss function ( Hadsell et al. , 2006 ) . The contrastive loss is defined such that when minimized , the representations of similar input points are pulled towards each other , while those of dissimilar input points are pushed away from each other . The contrastive loss has its roots in linear discriminant analysis ( Fisher , 1936 ) and is closely related to the triplet loss ( Weinberger et al. , 2006 ) . Recent approaches in contrastive learning are characterized by the InfoNCE loss proposed by Oord et al . ( 2018 ) . CPC uses InfoNCE as a lower bound of mutual information ( MI ) and maximizes this lowerbound , by using negative examples . Deep InfoMax ( Hjelm et al. , 2018 ) similarly proposes to use the idea of maximizing MI while considering global and local representations of an image . Hjelm et al . ( 2018 ) tested three bounds on MI : Donsker-Varadhan ( Donsker & Varadhan , 1983 . ) , Jensen-Shannon ( Nowozin et al. , 2016 ) , and InfoNCE ( Oord et al. , 2018 ) , and found that the InfoNCE objective resulted in the best downstream classification accuracies . Since then , several more advances in contrastive self-supervised learning have happened , such as AMDIM ( Bachman et al. , 2019 ) and CMC ( Tian et al. , 2019 ) , both of which focus on using multiple views of each image . Hénaff et al . ( 2019 ) extend CPC with an image patch prediction task , and YADIM ( Falcon & Cho , 2020 ) combines these ideas of augmentation and InfoNCE loss from both CPCv2 ( Hénaff et al. , 2019 ) and AMDIM ( Bachman et al. , 2019 ) under a single framework . The success of contrastive learning comes from using a large number of negative examples . Misra & van der Maaten ( 2020 ) empirically demonstrate with PIRL the benefits of using a large number of negative examples for downstream task performance . PIRL uses a momentum-updated memory bank ( Wu et al. , 2018 ) to provide this large cache of negatives . Memory bank models ( Wu et al. , 2018 ; Misra & van der Maaten , 2020 ) need to store and update representations for each data point and hence can not be scaled up efficiently . To remove the dependence on memory bank , MoCo ( He et al. , 2020 ; Chen et al. , 2020b ) instead introduces a momentum-updated encoder and a comparatively smaller queue of representations to set up positive and negative pairs for contrastive learning . SimCLR ( Chen et al. , 2020a ) removes memory banks and momentum-updated encoders and scales up the batch size to provide a large number of negatives from within each mini-batch . The necessity of a large quantity of negatives for the contrastive loss function to work well proves to be a major challenge in scaling up these methods .

Summary: The paper presents a third family for self-supervised learning that relies on generative modeling, in particular variational autoencoders (VAE). The proposed model uses data augmentation in the data space as a way to promote learning richer latent representations, where the augmentation strategy reflects data/domain-specific invariances and equivariances. This augmentation strategy regularizes the VAE  learning process and is used to replace the domain-agnostic KL divergence regularization. Learned representations are evaluated on image classification downstream tasks and compared with recent contrastive and non-contrastive self-supervised methods.