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: This paper presents a modification to the VAE training scheme. Here, the training objective is modified with the aim of encouraging the model to correctly learn invariances and equivariances of the data. Finally, the authors empirically evaluate the proposed model on standard datasets such as CIFAR-10, STL-10, and ImageNet.