Improving Transferability of Representations via Augmentation-Aware Self-Supervision
Year: 2021

Abstract:
Recent unsupervised representation learning methods have shown to be effective in a range of vision tasks by learning representations invariant to data augmentations such as random cropping and color jittering. However, such invariance could be harmful to downstream tasks if they rely on the characteristics of the data augmentations, e.g., locationor color-sensitive. This is not an issue just for unsupervised learning; we found that this occurs even in supervised learning because it also learns to predict the same label for all augmented samples of an instance. To avoid such failures and obtain more generalizable representations, we suggest to optimize an auxiliary self-supervised loss, coined AugSelf, that learns the difference of augmentation parameters (e.g., cropping positions, color adjustment intensities) between two randomly augmented samples. Our intuition is that AugSelf encourages to preserve augmentation-aware information in learned representations, which could be beneficial for their transferability. Furthermore, AugSelf can easily be incorporated into recent state-of-the-art representation learning methods with a negligible additional training cost. Extensive experiments demonstrate that our simple idea consistently improves the transferability of representations learned by supervised and unsupervised methods in various transfer learning scenarios. The code is available at https://github.com/hankook/AugSelf.

heading: 1 Introduction
text: 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.

heading: 2 Preliminaries: Augmentation-invariant representation learning
text: 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.

heading: 3 Auxiliary augmentation-aware self-supervision
text: 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 .

heading: 4 Experiments
text: Setup. We pretrain the standard ResNet-18 [20] and ResNet-50 on STL10 [23] and ImageNet1003 [31, 32], respectively. We use two recent unsupervised representation learning methods as baselines for pretraining: a contrastive method, MoCo v2 [1, 14], and a non-contrastive method, SimSiam [13]. For STL10 and ImageNet100, we pretrain networks for 200 and 500 epochs with a batch size of 256, respectively. For supervised pretraining, we pretrain ResNet-50 for 100 epochs with a batch size of 128 on ImageNet100.4 For augmentations, we use random cropping, flipping, color jittering, grayscaling, and Gaussian blurring following Chen and He [13]. In this section, our AugSelf predicts random cropping and color jittering parameters, i.e., AAugSelf = {crop, color}, unless otherwise stated. We set λ = 1.0 for STL10 and λ = 0.5 for ImageNet100. The other details and the sensitivity analysis to the hyperparameter λ are provided in the supplementary material. For ablation study (Section 4.2), we only use STL10-pretrained models.

heading: 4.1 Main results
text: Linear evaluation in various downstream tasks. We evaluate the pretrained networks in downstream classification tasks on 11 datasets: CIFAR10/100 [29], Food [24], MIT67 [36], Pets [37], Flowers [21], Caltech101 [38], Cars [39], Aircraft [40], DTD [41], and SUN397 [42]. They contain roughly 1k∼70k training images. We follow the linear evaluation protocol [25]. The detailed information of datasets and experimental settings is described in the supplementary material. Table 1 shows the transfer learning results in the various downstream tasks. Our AugSelf consistently improves (a) the recent unsupervised representation learning methods, SimSiam [13] and MoCo [13], in 10 out of 11 downstream tasks; and (b) supervised pretraining in 9 out of 11 downstream tasks. These consistent improvements imply that our method encourages to learn more generalizable representations.
Few-shot classification. We also evaluate the pretrained networks on various few-shot learning benchmarks: FC100 [33], Caltech-UCSD Birds (CUB200) [43], and Plant Disease [35]. Note that CUB200 and Plant Disease benchmarks require low-level features such as color information of birds and leaves, respectively, to detect their fine-grained labels. They are widely used in cross-domain few-shot settings [44, 45]. For few-shot learning, we perform logistic regression using the frozen representations f(x) without fine-tuning. Table 2 shows the few-shot learning performance of 5-way 1-shot and 5-way 5-shot tasks. As shown in the table, our AugSelf improves the performance of SimSiam [13] and MoCo [1] in all cases with a large margin. For example, for plant disease detection
3ImageNet100 is a 100-category subset of ImageNet [31]. We use the same split following Tian et al. [32]. 4We do not experiment supervised pretraining on STL10, as it has only 5k labeled training samples, which is
not enough for pretraining a good representation.
[35], we obtain up to 6.07% accuracy gain in 5-way 1-shot tasks. These results show that our method is also effective in such transfer learning scenarios.
Comparison with LooC. Recently, Xiao et al. [16] propose LooC that learns augmentation-aware representations via multiple augmentation-specific contrastive learning objectives. Table 3 shows head-to-head comparisons under the same evaluation setup following Xiao et al. [16].5 As shown in the table, our AugSelf has two advantages over LooC: (a) AugSelf requires the same number of augmented samples compared to the baseline unsupervised representation learning methods while LooC requires more, such that AugSelf does not increase the computational cost; (b) AugSelf can be incorporated with non-contrastive methods e.g., SimSiam [13], and SimSiam with AugSelf outperforms LooC in all cases.
Object localization. We also evaluate representations in an object localization task (i.e., bounding box prediction) that requires positional information. We experiment on CUB200 [46] and solve linear regression using representations pretrained by SimSiam [13] without or with our method. Table 4 reports `2 errors of bounding box predictions and Figure 4 shows the examples of the predictions. These results demonstrate that AugSelf is capable of learning positional information.
Retrieval. Figure 5 shows the retrieval results using pretrained models. For this experiment, we use the Flowers [21] and Cars [39] datasets and find top-4 nearest neighbors based on the cosine
5Since LooC’s code is currently not publicly available, we reproduced the MoCo baseline as reported in the sixth row in Table 3: we obtained the same ImageNet100 result, but different ones for CUB200 and Flowers.
similarity between representations f(x) where f is the pretrained ResNet-50 on ImageNet100. As shown in the figure, the representations learned by AugSelf are more color-sensitive.

heading: 4.2 Ablation study
text: Effect of augmentation prediction tasks. We first evaluate the proposed augmentation prediction tasks one by one without incorporating invariance-learning methods. More specifically, we pretrain fθ using only Laug for each aug ∈ {crop, flip, color, blur}. Remark that training objectives are different but we use the same set of augmentations. Table 5 shows the transfer learning results in various downstream tasks. We observe that solving horizontal flipping and Gaussian blurring prediction tasks results in worse or similar performance to a random initialized network in various downstream tasks, i.e., the augmentations do not contain task-relevant information. However, solving random cropping and color jittering prediction tasks significantly outperforms the random initialization in all downstream tasks. Furthermore, surprisingly, the color jittering prediction task achieves competitive performance in the Flowers [21] dataset compared to a recent state-of-the-art method, SimSiam [13]. These results show that augmentation-aware information are task-relevant and learning such information could be important in downstream tasks.
Based on the above observations, we incorporate random cropping and color jittering prediction tasks into SimSiam [13] when pretraining. More specifically, we optimize LSimSiam + λcropLcrop + λcolorLcolor where λcrop, λcolor ∈ {0, 1}. The transfer learning results are reported in Table 6. As shown in the table, each self-supervision task improves SimSiam consistently (and often significantly) in various downstream tasks. For example, the color jittering prediction task improves SimSiam by 6.33% and 11.89% in Food [24] and Flowers [21] benchmarks, respectively. When incorporating both tasks simultaneously, we achieve further improvements in almost all the downstream tasks. Furthermore, as shown in Figure 2, our AugSelf preserves augmentation-aware information as much as possible; hence our gain is consistent regardless of the strength of color jittering augmentation.
Different augmentations. We confirm that our method can allow to use other strong augmentations: rotation, which rotates an image by 0◦, 90◦, 180◦, 270◦ degrees randomly; and solarization, which inverts each pixel value when the value is larger than a randomly sampled threshold. Based on the default augmentation setting, i.e., AAugSelf = {crop, color}, we additionally apply each
augmentation with a probability of 0.5. We also evaluate the effectiveness of augmentation prediction tasks for rotation and solarization. Note that we formulate the rotation prediction as a 4-way classification task (i.e., ωrotdiff ∈ {0, 1, 2, 3}) and the solarization prediction as a regression task (i.e., ωsoldiff ∈ [−1, 1]). As shown in Table 7, we obtain consistent gains across various downstream tasks even if stronger augmentations are applied. Furthermore, in the case of rotation, we observe that our augmentation prediction task tries to prevent the performance degradation from learning invariance to rotations. For example, in CIFAR100 [29], the baseline loses 4.66% accuracy (54.90%→50.24%) when using rotations, but ours does only 0.37% (58.65%→58.28%). These results show that our AugSelf is less sensitive to the choice of augmentations. We believe that this robustness would be useful in future research on representation learning with strong augmentations.
Solving geometric and color-related pretext tasks. To validate that our AugSelf is capable of learning augmentationaware information, we try to solve two pretext tasks requiring the information: 4-way rotation (0◦, 90◦, 180◦, 270◦) and 6-way color channel permutation (RGB, RBG, . . ., BGR) classification tasks. We note that the baseline (SimSiam) and our method (SimSiam+AugSelf) do not observe rotated or
color-permuted samples in the pretraining phase. We train a linear classifier on top of pretrained representation without finetuning for each task. As reported in Table 8, our AugSelf solves the pretext tasks well even without their prior knowledge in pretraining; these results validate that our method learns augmentation-aware information.
Compatibility with other methods. While we mainly focus on SimSiam [13] and MoCo [1] in the previous section, our AugSelf can be incorporated into other unsupervised learning methods, SimCLR [2], BYOL [12], and SwAV [11]. Table 9 shows the consistent and significant gains by AugSelf across all methods and downstream tasks.

heading: 5 Related work
text: Self-supervised pretext tasks. For visual representation learning without labels, various pretext tasks have been proposed in literature [17, 18, 19, 47, 48, 49, 50, 51] by constructing self-supervision from an image. For example, Doersch et al. [17], Noroozi and Favaro [18] split the original image x into 3 × 3 patches and then learn visual representations by predicting relations between the patch locations. Instead, Zhang et al. [19], Larsson et al. [48] construct color prediction tasks by converting colorful images to gray ones. Zhang et al. [51] propose a similar task requiring to predict one subset of channels (e.g., depth) from another (e.g., RGB values). Meanwhile, Gidaris et al. [47], Qi et al.
[49], Zhang et al. [50] show that solving affine transformation (e.g., rotation) prediction tasks can learn high-level representations. These approaches often require specific preprocessing procedures (e.g., 3× 3 patches [17, 18], or specific affine transformations [49, 50]). In contrast, our AugSelf can be working with common augmentations such as random cropping and color jittering. This advantage allows us to incorporate AugSelf to the recent state-of-the-art frameworks like SimCLR [2] while not increasing the computational cost. Furthermore, we emphasize that our contribution is not only to construct AugSelf, but also finding on the importance of learning augmentation-aware representations together with the existing augmentation-invariant approaches.
Augmentations for unsupervised representation learning. Chen et al. [2], Tian et al. [15] found that the choice of augmentations plays a critical role in contrastive learning. Based on this finding, many unsupervised learning methods [1, 2, 11, 12, 13] have used similar augmentations (e.g., cropping and color jittering) and then achieved outstanding performance in ImageNet [31]. Tian et al. [15] discuss a similar observation to ours that the optimal choice of augmentations is task-dependent, but they focus on finding the optimal choice for a specific downstream task. However, in the pretraining phase, prior knowledge of downstream tasks could be not available. In this case, we need to preserve augmentation-aware information in representations as much as possible for unknown downstream tasks. Recently, Xiao et al. [16] propose a contrastive method for learning augmentation-aware representations. This method requires additional augmented samples for each augmentation; hence, the training cost increases with respect to the number of augmentations. Furthermore, the method is specialized to contrastive learning only, and is not attractive to be used for non-contrastive methods like BYOL [12]. In contrast, our AugSelf does not suffer from these issues as shown in Table 3 and 9.

heading: 6 Discussion and conclusion
text: To improve the transferability of representations, we propose AugSelf, an auxiliary augmentationaware self-supervision method that encourages representations to contain augmentation-aware information that could be useful in downstream tasks. Our idea is to learn to predict the difference between augmentation parameters of two randomly augmented views. Through extensive experiments, we demonstrate the effectiveness of AugSelf in various transfer learning scenarios. We believe that our work would guide many research directions for unsupervised representation learning and transfer learning.
Limitations. Even though our method provides large gains when we apply it to popular data augmentations like random cropping, it might not be applicable to some specific data augmentation, where it is non-trivial to parameterize augmentation such as GAN-based one [52]. Hence, an interesting future direction would be to develop an augmentation-agnostic way that does not require to explicitly design augmentation-specific self-supervision. Even for this direction, we think that our idea of learning the relation between two augmented samples could be used, e.g., constructing contrastive pairs between the relations.
Negative societal impacts. Self-supervised training typically requires a huge training cost (e.g., training MoCo v2 [14] for 1000 epochs requires 11 days on 8 V100 GPUs), a large network capacity (e.g., GPT3 [5] requires 175 billion parameters), therefore it raises environmental issues such as carbon generation [53]. Hence, efficient training methods [54] or distilling knowledge to a smaller network [55] would be required to ameliorate such environmental problems.
Acknowledgments and disclosure of funding
This work was mainly supported by Samsung Electronics Co., Ltd (IO201211-08107-01) and partly supported by Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST)).

Review summaries:
This paper aims at improving the transferability of representations, especially self-supervised pre-trained representations. Motivated by the observation that augmentation-invariance may be harmful to downstream tasks, e.g., location- or color-sensitive ones, they proposed an auxiliary augmentation-aware self-supervision loss named AugSelf. AugSelf learns the difference of augmentation parameters (e.g., cropping positions, color adjustment intensities) between two randomly augmented samples. Experiments show that AugSelf helps for transfer learning tasks with two pre-trained models and a supervised learning model, as well as few-shot learning setups. 

This paper proposes a simple regularization for self-supervised/supervised representation learning, which aims to benefit downstream tasks. Main idea is to make the network to be variant to different augmentation. Extensive experimental results are reported to show the effectiveness.

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.