Learning robust visual representations using data augmentation invariance

Keywords : deep neural networks ; visual cortex ; invariance ; data augmentation 1 INTRODUCTION . Deep artificial neural networks ( DNNs ) have borrowed much inspiration from neuroscience and are , at the same time , the current best model class for predicting neural responses across the visual system in the brain ( Kietzmann et al. , 2017 ; Kubilius et al. , 2018 ) . Yet , despite consensus about the benefits of a closer integration of deep learning and neuroscience ( Bengio et al. , 2015 ; Marblestone et al. , 2016 ) , important differences remain . Here , we investigate a representational property that is well established in the neuroscience literature on the primate visual system : the increasing robustness of neural responses to identity-preserving image transformations . While early areas of the ventral stream ( V1-V2 ) are strongly affected by variation in e.g . object size , position or illumination , later levels of processing are increasingly robust to such changes , as measured first in single neurons of the inferior temporal ( IT ) cortex of macaques Booth & Rolls ( 1998 ) and then in humans ’ ( Quiroga et al. , 2005 ; Isik et al. , 2013 ) . The cascaded achievement of invariance to such identity-preserving transformations has been proposed as a key mechanism for robust object recognition ( DiCarlo & Cox , 2007 ; Tacchetti et al. , 2018 ) . Learning such invariant representations has been a desired objective since the early days of artificial neural networks ( Simard et al. , 1992 ) . Accordingly , a myriad of techniques have been proposed to attempt to achieve tolerance to different types of transformations ( see Cohen & Welling ( 2016 ) for a review ) . Interestingly , recent theoretical work ( Achille & Soatto , 2018 ) has shown that invariance to “ nuisance factors ” should naturally emerge from the optimization process of deep models that minimize the information of the representations about the inputs , while retaining the minimum information about the target , as proposed by Tishby & Zaslavsky ( 2015 ) in the information bottleneck principle . Nevertheless , DNNs are still not robust to identity-preserving transformations , including simple image translations ( Zhang , 2019 ) . A remarkable extreme example are adversarial attacks ( Szegedy et al. , 2013 ) , in which small changes , imperceptible to the human brain , can alter the classification output of the network . Extending this line of research , we use the data augmentation framework ( Hernández-García & König , 2018 ) to show that DNNs , despite being trained on augmented data , are not robust to such transformations and the learned representations are even less disentangled than in the input space . This is likely related to the growing evidence that DNNs may exploit highly discriminative features that do not match human perception ( Jo & Bengio , 2017 ; Ilyas et al. , 2019 ; Wang et al. , 2019 ) . We postulate that this is due to the combination of their large capacity and the highly unconstrained learning setup of typical supervised deep models , and that incorporating elements from human visual perception and biological constraints can add a positive inductive bias that may yield better , more robust representations . Hence , inspired by the increasing invariance observed along the primate ventral visual stream , we subsequently propose data augmentation invariance , a simple , yet effective and efficient mechanism to improve the robustness of the representations : we include an additional , unsupervised term in the objective function that encourages the similarity between augmented examples within each batch . In sum , the main contributions of this paper are the following : • We propose a data augmentation invariance score that intuitively evaluates the robustness of the intermediate features of a neural network towards transformations of the input images . • We use this score to show that the intermediate representations of several popular architectures do not become more disentangled than in the input space , after training . • We propose an unsupervised , layer-wise learning objective that encourages the representational similarity between transformed examples . 2 METHODS . This section presents the procedure to empirically measure the robustness of the representations of a convolutional neural network and our proposal to improve the invariance . 2.1 EVALUATION OF INVARIANCE . To measure the invariance of the learned features under the influence of identity-preserving image transformations we compare the activations of a given image with the activations of a data augmented version of the same image . Consider the activations of an input image x at layer l of a neural network , which can be described by a function f ( l ) ( x ) ∈ RD ( l ) . We can define the distance between the activations of two input images xi and xj by their mean square difference : d ( l ) ( xi , xj ) = 1 D ( l ) D ( l ) ∑ k=1 ( f ( l ) k ( xi ) − f ( l ) k ( xj ) ) 2 ( 1 ) Following this , we compute the mean square difference between f ( l ) ( xi ) and a random transformation of xi , that is d ( l ) ( xi , G ( xi ) ) . G ( x ) refers to the stochastic function that transforms the input images according to a pre-defined data augmentation scheme . In order to assess the similarity between the activations of an image xi and of its augmented versions G ( xi ) we normalize it by the average similarity in the ( test ) set . We define the invariance score S ( l ) i of the transformation G ( xi ) at layer l of a model , with respect to a data set of size N , as follows : S ( l ) i = 1− d ( l ) ( xi , G ( xi ) ) 1 N ∑N j=1 d ( l ) ( xi , xj ) ( 2 ) Note that the invariance S ( l ) i takes the maximum value of 1 if the activations of xi and its transformed version G ( xi ) are identical . 2.2 DATA AUGMENTATION INVARIANCE . Most DNNs trained for object categorization are optimized through mini-batch gradient descent ( SGD ) , that is the weights are updated iteratively by computing the loss of a batch B of examples , instead of the whole data set at once . The models are typically trained for a number of epochs , E , which is a whole pass through the entire training data set of size N . That is , the weights are updated K = N|B| times each epoch . Data augmentation introduces variability into the process by performing a different , stochastic transformation of the data every time an example is fed into the network . However , with standard data augmentation , the model has no information about the identity of the images , that is , that different augmented examples , seen at different epochs , separated by N|B| iterations on average , correspond to the same seed data point . This information is potentially valuable and useful to learn better representations . For example , in biological vision , the high temporal correlation of the stimuli that reach the visual cortex may play a crucial role in the creation of robust connections ( Wyss et al. , 2006 ) . However , this is generally not exploited in supervised settings . In semi-supervised learning , where the focus is on learning from fewer labeled examples , data augmentation has been used as a source of variability together with dropout or random pooling , among others ( Laine & Aila , 2016 ) . In order to make use of this information and improve the robustness , we first propose to perform in-batch data augmentation by constructing the batches with M transformations of each example ( see Hoffer et al . ( 2019 ) for a similar idea ) . Additionally , we propose to modify the loss function to include an additional term that accounts for the invariance of the feature maps across multiple image samples . Considering the difference between the activations at layer l of two images , d ( l ) ( xi , xj ) , defined in Equation 1 , we define the data augmentation invariance loss at layer l for a given batch B as follows : L ( l ) inv = ∑ k 1 |Sk|2 ∑ xi , xj∈Sk d ( l ) ( xi , xj ) 1 |B|2 ∑ xi , xj∈B d ( l ) ( xi , xj ) ( 3 ) where Sk is the set of samples in the batch B that are augmented versions of the same seed sample xk . This loss term intuitively represents the average difference of the activations between the sample pairs that correspond to the same source image , relative to the average difference of all pairs . A convenient property of this definition is that Linv does not depend on the batch size nor the number of in-batch augmentations M = |Sk| . Furthermore , it can be efficiently implemented using matrix operations . Since both , certain representational invariance at L layers of the network and high object recognition performance at the network output are desired , we define the total loss as follows : L = ( 1− α ) Lobj + L∑ l=1 α ( l ) L ( l ) inv ( 4 ) where ∑L l=1 α ( l ) = α and Lobj is the loss associated with the object recognition objective , typically the cross-entropy between the object labels and the output of a softmax layer . All the results we report in this paper have been obtained by setting α = 0.1 and distributing the coefficients across the layers according to an exponential law , such that α ( l=L ) = 10α ( l=1 ) . This aims at simulating a probable response along the ventral visual stream , where higher regions are more invariant than the early visual cortex1 . 2.3 ARCHITECTURES AND DATA SETS . As test beds for our hypotheses and proposal we train three neural network architectures : all convolutional network , All-CNN-C ( Springenberg et al. , 2014 ) ; wide residual network , WRN-28-10 ( Zagoruyko & Komodakis , 2016 ) ; and DenseNet-BC ( Huang et al. , 2017 ) . All three architectures have been widely used in the deep learning literature and they have distinctive architectural characteristics : only convolutional layers , residual blocks and dense connectivity , respectively . We train the three architectures on the highly benchmarked data set for object recognition CIFAR-10 ( Krizhevsky & Hinton , 2009 ) . Additionally , in order to test our proposal on higher resolution images and a larger number of classes , we also train All-CNN and WRN on the tiny ImageNet data set , a subset of ImageNet ( Russakovsky et al. , 2015 ) with 100,000 64x64 training images that belong to 200 classes . All models are trained using a data augmentation scheme that consists in affine transformations , contrast adjustment and brightness adjustment ( see the details in Appendix A ) For the models trained on CIFAR-10 , the training hyperparameters ( learning rate , number of epochs , etc . ) are set as in the original papers , except that , following the recommendation by Hernández-García & König ( 2018 ) we do not use explicit regularization ( weight decay and dropout ) since comparable performance is obtained without them if data augmentation is used . On tiny ImageNet , All-CNN has three additional layers and is trained for 150 epochs , with batch size 128 and initial learning rate 0.01 decayed by 0.1 at epochs 100 and 125 . WRN is trained for 50 epochs , with batch size 32 and initial learning rate 0.01 decayed by 0.2 at epochs 30 and 40 . Both architectures are trained with stochastic gradient descent with Nesterov momentum 0.9 . Note that the hyperparameters were fine tuned for training only with the standard categorical crossentropy and with standard epoch-wise data augmentation . Therefore , they are likely suboptimal for our proposed data augmentation invariance . However , our aim is not achieving the best possible classification performance , but rather demonstrate the suitability of data augmentation invariance and analyze the learned representations . The invariance loss defined in Equation 3 was computed after the ReLU activation of each convolutional layer for All-CNN ; at the output of each residual block for WRN , and after the first convolution and the output of each dense block for DenseNet .

Summary: The paper proposes to explicitly improve the robustness of image-classification models to invariant transformations, via a secondary multi-task objective. The idea is that the secondary objective makes intermediate-layer representations invariant to transformations of the image that should lead to the same classification. The paper also establishes that the typical models do not actually learn such representations by themselves.