Distribution Estimation to Automate Transformation Policies for Self-Supervision

Our framework to estimate the distribution of visual transformations consists of two networks: a generator and a discriminator. The generator is constructed with three linear layers with 128 dimensional output except that the last layer is 6 dimensional output. We used LeakyReLU with $0.2$ slope for the first two layers, and the last layer has the tangent hyperbolic function to produce values between $-1$ and $1$. The generator takes 10 dimensional inputs sampled from Gaussian distribution, and outputs 6 parameters for affine transformation (rotation, scale, and translation). We found that the affine transformation parameters can be accurately estimated with two generators: i.e., the first one for the scale parameter and the other one for remaining five parameters. Similar with the affine case, the color transformation can be generated from 10 dimensional inputs (See the details in Appendix C). In this case, the generator outputs four dimensional parameters: brightness, saturation, contrast, and hue. And, the network architecture is the same with the affine transformation case, but the output dimension is four. We adopted a simple discriminator that has three linear layers with the LeakyReLU function. The input size of the first layer depends on the image size, and the output dimensions of the three layers are respectively 512, 256, and 1. The discriminator output represents a fake or a real, i.e., an image from the given dataset (real) or a transformed image using the generator (fake).

We trained these networks with the aid of WGAN-GP [31]. We adopted the basic settings of WGAN-GP: $\lambda=10$, $n_{critic}=5$, $\alpha=0.0001$, $\beta_{1}=0$, and $\beta_{2}=0.9$. With the batch size of $10$ and learning rate of $5e-5$, we trained the network for $500,000$ iterations, which took $6$ hours in our single NVIDIA TITAN Xp GPU. Fig. 4 illustrates three histograms obtained at different iteration points, and we observed that there is no big difference after 500,000 iterations.

Through adversarial learning, the discriminator is trained to distinguish the original images in the dataset from the images transformed by the generator given the reference subset. Our objective is to train the generator to deceive the discriminator such that the transformed images have a similar distribution with the original images. However, the discriminator may easily identify the generated images when they have artifacts like zero padding. These artifacts cause the network to learn an undesired property. For example, assume that the input image contains a dog of 45 degree rotated, and the objective is to find the amount of the rotation by generating a similar image using an upright dog image. However, when the image is rotated in the generator, the four corners are zero-padded, and hence it provides a crucial hint for the discriminator. Often, it leads to converge to a meaningless distribution which is not our objective. To mitigate this problem, we perform transformations first which is followed by a center-crop. For example, given a 32x32 image, a transformation is applied, and then the 24x24 image is obtained by the center-crop, which can cover 8 pixel shifts with no zero padding. For 256x256 image, the resulting image size is 196x196 where these two numbers were empirically selected through multiple trials. With this procedure, we can accurately estimate the distribution of visual transformations without an artifact effect.

We used the NIN architecture [32] as the backbone feature extraction network to learn the representations for downstream tasks. The NIN consists of four convolutional blocks, and each block contains three convolutional layers. Convolution, batch normalization, and ReLU are sequentially placed in a single convolutional layer. In training a pretext task for rotation, translation, and scaling categories, the classifier of a single linear layer is added after global average pooling. We used SGD with the batch size of $128$, momentum of $0.9$, weight decay of $5e-4$, and learning rate of $0.1$. We dropped the learning rates by a factor of $5$ after epochs $60$, $120$, and $160$. We trained the network in total for $200$ epochs. For a fair comparison among several pretext tasks, we used the same training procedure.

To evaluate the representation learned by self-supervision, a classifier is trained using the learned features. Specifically, in our experiments on Fashion MNIST [18], SVHN [19], CIFAR-10 [20], and CIFAR-100 [20], we followed the existing evaluation protocols [2, 13] by adding a classifier on top of the second convolutional block. The classifier consists of a single conv block, global average pooling, and one linear layer. Here, the conv block is the same with the conv3 in NIN architecture. We used SGD with batch size $128$, momentum $0.9$, weight decay $5e-4$, and learning rate of $0.1$, which is dropped by a factor of $5$ after epochs $35$, $70$, and $85$. We trained the network in total for $100$ epochs.

We introduce four pretext tasks, R, T, S, and RST, which indicate rotation, translation, scaling, and the combination of all, respectively. For the rotation case, the unit is the rotation degree. Translation and scaling units are the ratio of shift amount to the image size and the scaling ratio. Note that the translation task involves 5-way classification: left and right shift along with $x$ and $y$ axis, and 0 translation.

When constructing a pretext task of rotation, translation, and scaling, we also encounter the same artifact problem. The artifact induced by the transformation for a pretext task would be a big hint for the network, and hence the useful representation for downstream tasks may not be obtained. With this reason, we consider the pretext tasks of each transformation as follows: 1) For rotation, we use 90, 180, 270, and 360 degree which do not incur the zero-padding, 2) For translation and scaling, we perform the center crop depending on the task parameter: e.g. when we define the translation task of 3 pixels in 32x32 image, we crop the center image of 26x26.

We describe how to parameterize a distribution over the set of visual transformations, particularly affine transformations. This discussion can be extended to projective and spline transformations as well. For the affine transformation, we include translation, scaling, similarity, reflection, rotation, and shear. The affine matrix is parameterized as follows:

$$t=\begin{bmatrix}a_{11}&a_{12}&a_{13}\\ a_{21}&a_{22}&a_{23}\\ 0&0&1\end{bmatrix},\quad\text{where}\quad a_{11},a_{12},...,a_{23}\in\theta.$$

(1)

These parameters are sampled from $p_{\theta}$, and this matrix is applied to the coordinate system. We adapt Spatial Transformer Network [33] allowing loss gradients to flow back to the sampling grid coordinates as well as to the input image. For example, in case of affine transformation, the sampling grid is result of warping the regular grid with an affine matrix, as in (1). If we use a bilinear sampling kernel to map the output pixel value, the output image $V$ is obtained as follows:

$$O_{i}^{c}=\sum_{n}^{H}\sum_{m}^{W}I_{nm}^{c}\text{max}(0,1-\left|u_{i}^{s}-m\right|)\text{max}(0,1-\left|v_{i}^{s}-n\right|),$$

(2)

where $O_{i}^{c}$ indicates the output value for pixel $i$ at location $(u_{i}^{t},v_{i}^{t})$ of channel $c$, and $(u_{i}^{t},v_{i}^{t})$ are the target coordinates of the regular grid. $I_{nm}^{c}$ is the value at location $(n,m)$ of the input image with height $H$ and width $W$, and $(u_{i}^{s},v_{i}^{s})$ indicates the spatial location of the input corresponding to the output coordinate through the sampling grid. We can define the gradients with respect to $I_{nm}^{c}$ and $(u_{i}^{s},v_{i}^{s})$ such that the backpropagation of the loss can be flowed through this sampling mechanism.

In all previous experiments, we considered the transformations of rotation, scale, and translation. This section will describe how we can parameterize the color transformation. Specifically, we consider changing the brightness, saturation, contrast, and hue of an image. We introduce learnable parameters for these transformations and train the mapping network such that each parameter represents the transformation distribution of the dataset.

First, the brightness change can be simply expressed as follows:

$$x_{brt}=x\alpha_{brt},$$

(3)

where $\alpha_{brt}$ is a learnable scale factor. Second, we can change the saturation of $x$ by using a linear combination of the original image $x$ and the gray-scaled version of the image $x_{gray}$ as follows:

$$x_{sat}=x\alpha_{sat}+x_{gray}(1-\alpha_{sat}),\quad\text{{\color[rgb]{0,0,0}and}}\quad x_{gray}=0.299x_{r}+0.587x_{g}+0.114x_{b},$$

(4)

where $x_{r}$, $x_{g}$, and $x_{b}$ indicate three color channels of $x$. The scale factor $\alpha_{sat}$ is in the range of $[0,1]$. And third, we can change the contrast of $x$ by calculating a linear combination of $x$ and the average of $x_{gray}$ over all spatial dimensions as follows:

$$x_{con}=x\alpha_{con}+mean(x_{gray})(1-\alpha_{con}),$$

(5)

where $\alpha_{con}$ is a scale parameter. Finally, we resort to a linear approximation for the hue change in the YIQ color space. We firstly convert RGB to YIQ, and then apply rotation to the IQ components. The transformation from RGB to YIQ is denoted by $T_{YIQ}$, and we can obtain the following expression:

$$\begin{bmatrix}Y\\ I\\ Q\end{bmatrix}=T_{YIQ}\begin{bmatrix}R\\ G\\ B\end{bmatrix},\quad\text{where}\quad T_{YIQ}=\begin{bmatrix}0.299&0.587&0.114\\ 0.596&-0.275&-0.321\\ 0.212&-0.523&0.311\end{bmatrix}.$$

(6)

In the YIQ color space, we can change the hue of an image by rotating the IQ components with a rotation matrix:

$$R_{\theta_{hue}}=\begin{bmatrix}1&0&0\\ 0&\text{cos}\theta_{hue}&-\text{sin}\theta_{hue}\\ 0&\text{sin}\theta_{hue}&\text{cos}\theta_{hue}\end{bmatrix},$$

(7)

where $\theta_{hue}=2\pi\alpha_{hue}$, and $\alpha_{hue}$ is a learnable parameter. In summary, the hue change can be expressed by the following matrix multiplication:

$$x_{hue}=T_{RGB}R_{\theta_{hue}}T_{YIQ}x,$$

(8)

where $T_{RGB}$ is the inverse matrix of $T_{YIQ}$. Using the procedure, we can employ learnable parameters for color transformation, and hence we can estimate the distribution of color transformations present in the dataset.

We plot the histogram of transformation parameters for Fashion MNIST [18] in Fig. 5. As illustrated in Fig. 5, the Fashion MNIST dataset consists of well-aligned images, i.e., almost no transformation is present. This suggests that the network can learn powerful representations with pretext tasks constructed by any visual transformation.

Furthermore, we illustrated the histograms of output values obtained from the mapping network on four visual recognition datasets, CropDisease [34], EuroSAT [35], ISIC2018 [36, 37], and ChestX [38].

Even when the images are big, we found that the mapping network can represent the transformation distribution well as shown in Fig 6. The EuroSAT and ISIC2018 dataset have the rotation transformation almost over the entire range, $[-\pi,\pi]$. Visually checking the image samples, we can find the two datasets already contain various rotated images.

To compare with RotNet [2] and VTSS [13], we set hyper-parameters with the same condition as others (using horizontal flip for data augmentation). We reported the accuracy in Table 1a.

We want to show that the pretext tasks selected by our estimated distribution can force the network to learn more meaningful representations than the manually selected one. In Table 1b,we define the pretext task, RST, which is the combined pretext task of three types of visual transformations. Baseline selects RST manually and ours defines RST based on the estimated distribution. We use rotation with 0, 90, 180, and 270, translation with 0.1, 0.1, and scaling x1, x1.14, x1.33 for the baseline RST while we define the pretext task carefully based on estimated distribution. Note that we do not use data augmentation technique to show the effectiveness of each pretext task separately.

We consider a contrastive learning framework based on SimCLR [4]. We observed that the contrastive loss with the augmented versions appropriate to the target dataset helps the network to learn generalizable representations even when the domain gap between source and target is large, and target data is scarce. Intuitively, this indicates proper pretext tasks for contrastive loss functions are effective mechanisms to derive meaningful representations for downstream tasks.

We apply the proposed automated transformation policy to generating augmented versions used for contrastive learning. As such, we retrain the base network, ResNet-50 [39], with the standard supervised loss on ImageNet [40], then fine-tune the base network, the added classifier, and the projector on the target dataset with both cross-entropy and contrastive losses. Specifically, we use supervised contrastive loss function [41] with the SimCLR [4] framework. It is to be noted that supervised contrastive loss function builds on the contrastive self-supervised literature by leveraging label information. We conduct experiments on cross-domain benchmark datasets (CropDisease [34], EuroSAT [35], ISIC2018 [36, 37], and ChestX [38]) where we split the total dataset such that $5\%$ was used for training while the remaining was used for testing.

We compare four methods: fine-tuning the network using (a) cross-entropy loss (Baseline) and with additional contrastive loss of (b) random rotation from -180 to 180 (SimCLR with Rot.), (c) random affine transformations, which contain rotation from -180 to 180, translation from -0.5 to 0.5, and scaling from 0.5 to 1.5 (SimCLR with Aff.) and (d) automated transformation policies within the same range of (c) and obtained from the target dataset (SimCLR with ATP). Table 3 shows the results of fine-tuning with the contrastive loss function. We found that some pretext tasks, i.e., data augmentations, are not helpful to the target tasks as pointed out by [42]. Our automated transformation policies for contrastive learning leads to better performance than the naive approach of using random augmentation. It shows that not only types of transformations, e.g., color or spatial transformation, but also the degree of transformations, e.g., 30 or 60 degrees rotation impact the performance significantly in contrastive learning.

In the future, we will conduct experiments with pretraining on a large-scale dataset and transfer to other datasets for evaluation. Since the large-scale dataset consists of images following a complex distribution, the manual selection of pretext tasks cannot be the best one and it may produce transformation conflict. Still, we believe that our policies will work well on a large-scale dataset.