Improving Transferability of Representations via Augmentation-Aware Self-Supervision

1 Introduction . Unsupervised representation learning has recently shown a remarkable success in various domains , e.g. , computer vision [ 1 , 2 , 3 ] , natural language [ 4 , 5 ] , code [ 6 ] , reinforcement learning [ 7 , 8 , 9 ] , and graphs [ 10 ] . The representations pretrained with a large number of unlabeled data have achieved outstanding performance in various downstream tasks , by either training task-specific layers on top of the model while freezing it or fine-tuning the entire model . In the vision domain , the recent state-of-the-art methods [ 1 , 2 , 11 , 12 , 13 ] learn representations to be invariant to a pre-defined set of augmentations . The choice of the augmentations plays a crucial role in representation learning [ 2 , 14 , 15 , 16 ] . A common choice is a combination of random cropping , horizontal flipping , color jittering , grayscaling , and Gaussian blurring . With this choice , learned representations are invariant to color and positional information in images ; in other words , the representations lose such information . On the contrary , there have also been attempts to learn representations by designing pretext tasks that keep such information in augmentations , e.g. , predicting positional relations between two patches of ∗Work done while at University of Michigan . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . an image [ 17 ] , solving jigsaw puzzles [ 18 ] , or predicting color information from a gray image [ 19 ] . These results show the importance of augmentation-specific information for representation learning , and inspire us to explore the following research questions : when is learning invariance to a given set of augmentations harmful to representation learning ? and , how to prevent the loss in the recent unsupervised learning methods ? Contribution . We first found that learning representations with an augmentation-invariant objective might hurt its performance in downstream tasks that rely on information related to the augmentations . For example , learning invariance against strong color augmentations forces the representations to contain less color information ( see Figure 2a ) . Hence , it degrades the performance of the representations in color-sensitive downstream tasks such as the Flowers classification task [ 21 ] ( see Figure 2b ) . To prevent this information loss and obtain more generalizable representations , we propose an auxiliary self-supervised loss , coined AugSelf , that learns the difference of augmentation parameters between the two augmented samples ( or views ) as shown in Figure 1 . For example , in the case of random cropping , AugSelf learns to predict the difference of cropping positions of two randomly cropped views . We found that AugSelf encourages the self-supervised representation learning methods , such as SimCLR [ 2 ] and SimSiam [ 13 ] , to preserve augmentation-aware information ( see Figure 2a ) that could be useful for downstream tasks . Furthermore , AugSelf can easily be incorporated into the recent unsupervised representation learning methods [ 1 , 2 , 12 , 13 ] with a negligible additional training cost , which is for training an auxiliary prediction head φ in Figure 1 . Somewhat interestingly , we also found that optimizing the auxiliary loss , AugSelf , can even improve the transferability of representations learned under the standard supervised representation learning scenarios [ 22 ] . This is because supervised learning also forces invariance , i.e. , assigns the same label , for all augmented samples ( of the same instance ) , and AugSelf can help to keep augmentation-aware knowledge in the learned representations . We demonstrate the effectiveness of AugSelf under extensive transfer learning experiments : AugSelf improves ( a ) two unsupervised representation learning methods , MoCo [ 1 ] and SimSiam [ 13 ] , in 20 of 22 tested scenarios ; ( b ) supervised pretraining in 9 of 11 ( see Table 1 ) . Furthermore , we found that AugSelf is also effective under few-shot learning setups ( see Table 2 ) . Remark that learning augmentation-invariant representations has been a common practice for both supervised and unsupervised representation learning frameworks , while the importance of augmentationawareness is less emphasized . We hope that our work could inspire researchers to rethink the under-explored aspect and provide a new angle in representation learning . 2 Preliminaries : Augmentation-invariant representation learning . In this section , we review the recent unsupervised representation learning methods [ 1 , 2 , 11 , 12 , 13 ] that learn representations by optimizing augmentation-invariant objectives . Formally , let x be an image , tω be an augmentation function parameterized by an augmentation parameter ω , v = tω ( x ) be the augmented sample ( or view ) of x by tω , and f be a CNN feature extractor , such as ResNet [ 20 ] . Generally speaking , the methods encourage the representations f ( v1 ) and f ( v2 ) to be invariant to the two randomly augmented views v1 = tω1 ( x ) and v2 = tω2 ( x ) , i.e. , f ( v1 ) ≈ f ( v2 ) for ω1 , ω2 ∼ Ω where Ω is a pre-defined augmentation parameter distribution . We now describe the recent methods one by one briefly . For simplicity , we here omit the projection MLP head g ( · ) which is widely used in the methods ( see Figure 1 ) . Instance contrastive learning approaches [ 1 , 2 , 26 ] minimize the distance between an anchor f ( tω1 ( x ) ) and its positive sample f ( tω2 ( x ) ) , while maximizing the distance between the anchor f ( tω1 ( x ) ) and its negative sample f ( tω3 ( x ′ ) ) . Since contrastive learning performance depends on the number of negative samples , a memory bank [ 26 ] , a large batch [ 2 ] , or a momentum network with a representation queue [ 1 ] has been utilized . Clustering approaches [ 11 , 27 , 28 ] encourage two representations f ( tω1 ( x ) ) and f ( tω2 ( x ) ) to be assigned into the same cluster , in other words , the distance between them will be minimized . Negative-free methods [ 12 , 13 ] learn to predict the representation f ( v1 ) of a view v1 = tω1 ( x ) from another view v2 = tω2 ( x ) . For example , SimSiam [ 13 ] minimizes ‖h ( f ( v2 ) ) − sg ( f ( v1 ) ) ‖22 where h is an MLP and sg is the stop-gradient operation . In these methods , if h is optimal , then h ( f ( v2 ) ) = Eω1∼Ω [ f ( v1 ) ] ; thus , the expectation of the objective can be rewritten as Varω∼Ω ( f ( v ) ) . Therefore , the methods can be considered as learning invariance with respect to the augmentations . Supervised learning approaches [ 22 ] also learn augmentation-invariant representations . Since they often maximize exp ( c > y f ( t ( x ) ) ) / ∑ y′ exp ( c > y′f ( t ( x ) ) ) where cy is the prototype vector of the label y , f ( t ( x ) ) is concentrated to cy , i.e. , cy ≈ f ( tω1 ( x ) ) ≈ f ( tω2 ( x ) ) . These approaches encourage representations f ( x ) to contain shared ( i.e. , augmentation-invariant ) information between tω1 ( x ) and tω2 ( x ) and discard other information [ 15 ] . For example , if tω changes color information , then to satisfy f ( tω1 ( x ) ) = f ( tω2 ( x ) ) for any ω1 , ω2 ∼ Ω , f ( x ) will be learned to contain no ( or less ) color information . To verify this , we pretrain ResNet-18 [ 20 ] on STL10 [ 29 ] using SimSiam [ 13 ] with varying the strength s of the color jittering augmentation . To measure the mutual information between representations and color information , we use the InfoNCE loss [ 30 ] . We here simply encode color information as RGB color histograms of an image . As shown in Figure 2a , using stronger color augmentations leads to color-relevant information loss . In classification on Flowers [ 21 ] and Food [ 24 ] , which is color-sensitive , the learned representations containing less color information result in lower performance as shown in Figure 2b and 2c , respectively . This observation emphasizes the importance of learning augmentation-aware information in transfer learning scenarios . 3 Auxiliary augmentation-aware self-supervision . In this section , we introduce auxiliary augmentation-aware self-supervision , coined AugSelf , which encourages to preserve augmentation-aware information for generalizable representation learning . To ! crop = ( ycenter , xcenter , H , W ) = ( 0.4 , 0.3 , 0.6 , 0.4 ) < latexit sha1_base64= '' c2I3TbFwd7fF3l8jpvjYz2q2zxw= '' > AAADCHicfVFda9swFJW9r9b7SrfHvngLKckwxm7TdTBSCttDX8Y6WJpClBlZuU5FZctI8mgw/gP7NXsbfd0f2PP+zeTELVtSdkHcwzlXV8fHcc6Z0kHw27Lv3L13/8HGpvPw0eMnT1tbz06VKCSFIRVcyLOYKOAsg6FmmsNZLoGkMYdRfPGu1kdfQSomss96nsMkJbOMJYwSbaio9auDRQozEoVfcEr0udZlwlleDUKnUW54KkVeReGgG/h9L/APvGXf6zm3De4MuvPohoFMg6y8yzXm2Bv1MHZ3rrfumfO63tzrePht4213xVtw7W1diVrtwA8W5a6DsAFt1NRJtGW9wFNBi9QYopwoNQ6DXE9KIjWjHCoHFwpyQi/IDMYGZiQFNSkXyVduxzBTNxHSnEy7C/bvGyVJlZqnsZmsfapVrSZv08aFTt5MSpblhYaMLh9KCu5q4da/0Z0yCVTzuQGESma8uvScSEJNqspx8HswHyPhg1n8MQdJtJCvSkzkLCWXVdn0/42xbDlmumNCDVcjXAenu37Y9/c/9dtHh028G2gbvURdFKIDdISO0QkaImrtW2NraoH9zf5u/7CvlqO21dx5jv4p++cfYObyNw== < /latexit > ! flip = [ v is flipped ] = 1 < latexit sha1_base64= '' TvSRxsaMNaXoNuiBbxIk0oE3xpE= '' > AAAES3icfVPbbtQwEHWzLZRwa+GRF4uqVYuiKGkXWmlVVAkk+oIoEr1I9bJynMmutc5FtlN1FeUTeIWv4gP4Dt4QDzgXlZJdYcmZyZnxmTPjJMgEV9rzfixZveWVO3dX79n3Hzx89Hht/cmZSnPJ4JSlIpUXAVUgeAKnmmsBF5kEGgcCzoPpmyp+fgVS8TT5pGcZDGM6TnjEGdUGGq1baJOkMYzpZxJTPdG6CEQuS3tz65BouNaF0qFLBiFcmWcakcE7mivFaUIGGE9BJiBKTEiV77ue3SWrBdZs20QYVSEd3dSRfDzRpdPFWZpoSdWCiKILwEkO5Q4huK7huXuOkeF47kFld2pljaSRfyMqEjwrD/15sTLNypFf0fQNxb7T2L2dxZl1xdlf3ZBokKVzPYccO+e3FNaUZr+q6GuFHfJanm3mX70HQVz45WXjR8VV2dwK5gpXeRmE5bDi2Kr6ccigbXa30+zNzcxHRmsbnuvVC887futsoHadjNaXdkmYsjw27TFBlbr0vUwPCyo1ZwJKm+QKMsqmdAyXxk1oDGrYfgp40yAhjlJpdqJxjd4+UdBYqVkcmMxKp+rGKnBhzEyqU1pHB8OCJ1muIWFN5SgXWKe4+hNwyCUwLWbGoUxyIx6zCZWUmUtTtk3egulOwntT6UMGkupUvigIleOYXpdFa/+XxpMmzVjbTNnvznTeOdt1/b778mN/4+h1O+9V9Aw9R9vIR/voCB2jE3SKmDW2vlhfrW+9772fvV+9302qtdSeeYr+WcsrfwDT1GPl < /latexit > ! color = ( bright , contrast , sat , hue ) = ( 0.3 , 1.0 , 0.8 , 1.0 ) < latexit sha1_base64= '' 7ytTg6EsrRlYxiLNDN7XjAGaf20= '' > AAAES3icfVPbbtQwEE2zLZRwa+GRF4tqqxZFUdIutNKqqBJI9AVRJHqR6mXlOJNda52LbKfqKvIn8ApfxQfwHbwhHnA2obTZFZacGZ05njkzdsKcM6l8/8eS3VleuXN39Z5z/8HDR4/X1p+cyqwQFE5oxjNxHhIJnKVwopjicJ4LIEnI4SycvKniZ5cgJMvST2qawyAho5TFjBJloOG6bXVxlsCIfMYJUWOlypAXQjvdzQOs4EqVUkUe7kdwab5ZjPvvSCElIynuIzQBkQLXCOOKH3i+08o106edzYMtzI2miAyvqwg2GivttnGapUoQuSAiyQJwXIDexhhVJXxv1zUaXN/br+z2TFYtaBhcS4o5y/VB4LTbpiLL9TCo0vRMij23trvbi5lVw1vTf7IhVSC0ezWHHLlnM4Hdzb+pd81+VaW/qfC2vtn4KyAMkzLQF7Ufl5e6vhTEJKqIOUR60Izfd7ou7jf97rT69a8nMRcZrm34nj9baN4JGmfDatbxcH1pB0cZLRLTIeVEyovAz9WgJEIxykE7uJCQEzohI7gwbkoSkIPmLaCuQSIUZ8LsVKEZevNESRIpp0lomJVO2Y5V4MKYGVWrtIr3ByVL80JBSuvKccGRylD1J6CICaCKT41DqGBGPKJjIgg19yYdB78F052A96bShxwEUZl4UWIiRgm50mVj/0djaU0z1jFTDtoznXdOd7yg57382Ns4fN3Me9V6Zj23tqzA2rMOrSPr2DqxqD2yv9hf7W+d752fnV+d3zXVXmrOPLVureWVP3ewY+Q= < /latexit > ! blur = std . dev . of Gaussian kernel = 1.0 < latexit sha1_base64= '' oc86yvF87jNFiI5wyqANY4moYgA= '' > AAAEVHicfVRfa9RAEE9zVduotdVHX4LHlVZCSK7VFo6TgoJ9ESvYP9A9j81mcrd0swm7m6NH2K/hq34qwe/ig5tcqO2luLCZYea3M7+Z2U2UMypVEPxesTurDx4+Wlt3Hj95uvFsc+v5mcwKQeCUZCwTFxGWwCiHU0UVg4tcAE4jBufR1fvKfz4DIWnGv6p5DqMUTzhNKMHKmMZb9jrKUpjgbyjFaqpUGbFCaGd7iBRcq1Kq2EeDGGbmmyVo8BEXUlLM0eAKBAemXYQMOPQDp7cUqOamnd72cAcxQyjG45scgk6mioOU2lv2kYwrgaVqeyRWhahpt33TAvQuQm6dLvD3PMPIC/zDSu5WHBt24/CGX8Jorodhm7fIcj0OqzD7JsSBt5B7u/cj64zzf/SBKxDau25Zjr3zWwzrkGa/rcLfZniXX4WuDVGUlqG+XOhJOdOL+bhUuhUwh1iP6ijbQzMKDw2aevtL9d7Mqe0Zb3YDP6iX21bCRulazToZb630UZyRIjUVEoalvAyDXI1KLBQlDLSDCgk5Jld4ApdG5TgFOWouhtszlthNMmE2V25tvX2ixKmU8zQyyIqnXPZVxnt9plVLqVVyOCopzwsFnCwyJwVzVeZWT8KNqQCi2NwomAhqyLtkigUmZm7ScdAHMNUJ+GQyfc7BXMFMvC4RFpMUX+uykf+DUb6AGemYLofLPW0rZ30/3PfffNnvHr1r+r1mvbReWTtWaB1YR9axdWKdWsTO7e/2D/tn51fnz6r5Syyg9kpz5oV1Z61u/AXc52a7 < /latexit > Original image Random cropping Color jittering Gaussian blurring Horizontal flipping Figure 3 : Examples of the commonly-used augmentations and their parameters ωaug . be specific , we add an auxiliary self-supervision loss , which learns to predict the difference between augmentation parameters of two randomly augmented views , into existing augmentation-invariant representation learning methods [ 1 , 2 , 11 , 12 , 13 ] . We first describe a general form of our auxiliary loss , and then specific forms for various augmentations . For conciseness , let θ be the collection of all parameters in the model . Since an augmentation function tω is typically a composition of different types of augmentations , the augmentation parameter ω can be written as ω = ( ωaug ) aug∈A where A is the set of augmentations used in pretraining ( e.g. , A = { crop , flip } ) , and ωaug is an augmentation-specific parameter ( e.g. , ωcrop decides how to crop an image ) . Then , given two randomly augmented views v1 = tω1 ( x ) and v2 = tω2 ( x ) , the AugSelf objective is as follows : LAugSelf ( x , ω1 , ω2 ; θ ) = ∑ aug∈AAugSelf Laug ( φaugθ ( fθ ( v1 ) , fθ ( v2 ) ) , ω aug diff ) , where AAugSelf ⊆ A is the set of augmentations for augmentation-aware learning , ωaugdiff is the difference between two augmentation-specific parameters ωaug1 and ω aug 2 , Laug is an augmentationspecific loss , and φaugθ is a 3-layer MLP for ω aug diff prediction . This design allows us to incorporate AugSelf into the recent state-of-the-art unsupervised learning methods [ 1 , 2 , 11 , 12 , 13 ] with a negligible additional training cost . For example , the objective of SimSiam [ 13 ] with AugSelf can be written as Ltotal ( x , ω1 , ω2 ; θ ) = LSimSiam ( x , ω1 , ω2 ; θ ) + λ · LAugSelf ( x , ω1 , ω2 ; θ ) , where λ is a hyperparameter for balancing losses . Remark that the total objective Ltotal encourages the shared representation f ( x ) to learn both augmentation-invariant and augmentation-aware features . Hence , the learned representation f ( x ) also can be useful in various downstream ( e.g. , augmentationsensitive ) tasks.2 In this paper , we mainly focus on the commonly-used augmentations in the recent unsupervised representation learning methods [ 1 , 2 , 11 , 12 , 13 ] : random cropping , random horizontal flipping , color jittering , and Gaussian blurring ; however , we remark that different types of augmentations can be incorporated into AugSelf ( see Section 4.2 ) . In the following , we elaborate on the details of ωaug and Laug for each augmentation . The examples of ωaug are illustrated in Figure 3 . Random cropping . The random cropping is the most popular augmentation in vision tasks . A cropping parameter ωcrop contains the center position and cropping size . We normalize the values by the height and width of the original image x , i.e. , ωcrop ∈ [ 0 , 1 ] 4 . Then , we use ` 2 loss for Lcrop and set ωcropdiff = ω crop 1 − ωcrop2 . Random horizontal flipping . A flipping parameter ωflip ∈ { 0 , 1 } indicates the image is horizontally flipped or not . Since it is discrete , we use the binary cross-entropy loss for Lflip and set ωflipdiff = 1 [ ωflip1 = ω flip 2 ] . Color jittering . The color jittering augmentation adjusts brightness , contrast , saturation , and hue of an input image in a random order . For each adjustment , its intensity is uniformly sampled from a pre-defined interval . We normalize all intensities into [ 0 , 1 ] , i.e. , ωcolor ∈ [ 0 , 1 ] 4 . Similarly to cropping , we use ` 2 loss for Lcolor and set ωcolordiff = ωcolor1 −ωcolor2 . 2We observe that our augmentation-aware objective LAugSelf does not interfere with learning the augmentationinvariant objective , e.g. , LSimSiam . This allows f ( x ) to learn augmentation-aware information with a negligible loss of augmentation-invariant information . A detailed discussion is provided in the supplementary material . Table 1 : Linear evaluation accuracy ( % ) of ResNet-50 [ 20 ] and ResNet-18 pretrained on ImageNet100 [ 31 , 32 ] and STL10 [ 23 ] , respectively . Bold entries are the best of each baseline . Method CIFAR10 CIFAR100 Food MIT67 Pets Flowers Caltech101 Cars Aircraft DTD SUN397 ImageNet100-pretrained ResNet-50 SimSiam 86.89 66.33 61.48 65.75 74.69 88.06 84.13 48.20 48.63 65.11 50.60 + AugSelf ( ours ) 88.80 70.27 65.63 67.76 76.34 90.70 85.30 47.52 49.76 67.29 52.28 MoCo v2 84.60 61.60 59.37 61.64 70.08 82.43 77.25 33.86 41.21 64.47 46.50 + AugSelf ( ours ) 85.26 63.90 60.78 63.36 73.46 85.70 78.93 37.35 39.47 66.22 48.52 Supervised 86.16 62.70 53.89 52.91 73.50 76.09 77.53 30.61 36.78 61.91 40.59 + AugSelf ( ours ) 86.06 63.77 55.84 54.63 74.81 78.22 77.47 31.26 38.02 62.07 41.49 STL10-pretrained ResNet-18 SimSiam 82.35 54.90 33.99 39.15 44.90 59.19 66.33 16.85 26.06 42.57 29.05 + AugSelf ( ours ) 82.76 58.65 41.58 45.67 48.42 72.18 72.75 21.17 33.17 47.02 34.14 MoCo v2 81.18 53.75 33.69 39.01 42.34 61.01 64.15 16.09 26.63 41.20 28.50 + AugSelf ( ours ) 82.45 57.17 36.91 41.67 43.80 66.96 66.02 17.53 28.02 45.21 30.93 Gaussian blurring . This blurring operation is widely used in unsupervised representation learning . The Gaussian filter is constructed by a single parameter , standard deviation σ = ωblur . We also normalize the parameter into [ 0 , 1 ] . Then , we use ` 2 loss for Lblur and set ωblurdiff = ωblur1 −ωblur2 .

Summary: This paper presents AugSelf, an auxiliary self-supervised task in contrastive learning to improves the transferability of learned representations. The task predicts the difference of augmentation parameters between two randomly augmented samples, and therefore forces the representation to preserve the information which could be potentially discarded by augmentations. Extensive experiments demonstrate the effectiveness of the proposed method.