Partitioning Image Representation in Contrastive Learning

In this section, we check qualitatively if the partitioned representation can split content and style features, especially on VAE. A traversal map of the latent space of VAE can show what features are learned by changing the value of only one dimension in the latent space. We used two datasets: Fashion MNIST [22] and colored MNIST where we add a feature about color on MNIST [23].

We train a VAE model with the partitioned representation and then visualize the latent space to understand what is represented in each dimension. The model has three convolution layers for the encoder and decoder, and 10 dimensions for the latent space with 7 for the content part and 3 for the style part. For each layer, the kernel size, stride, and padding size are 3, 2, and 1, respectively. Also, the channel size of the layers in the encoder are 32, 64, and 128, respectively (reverse order for the decoder). We use the Adam optimizer with the learning rate of 0.001. For contrastive learning, positive images are randomly selected from the same class as the anchor image.

We first train the model on the Fashion MNIST dataset, and Figure 4 shows traversal maps of the latent space. In each traversal map corresponding to one row, the center image is the reconstructed image from the latent space. Given the variance $\sigma_{i}$ for each dimension, we add $t\times\sigma_{i}$ with $t\in[-4,4]$ on only one dimension of the latent vector and then reconstruct the images from the manipulated latent vectors. As shown in Figure 4, the first seven dimensions (content part) show a variation from one class to another class, representing class-related features. However, the last three dimensions (style part) show variations within the class of the input image, which means they do not represent class-dependent features. Note that the common features between the anchor and positive images are class-related features because they are from the same class, otherwise, the unique features are about the style of each image.

For further investigation, we inject a content or style feature on purpose with colored MNIST. There are two cases of color injection: biased and unbiased cases as shown in Figure 5. The biased one is the case where each digit has its unique color, which means that color is a class-related feature shared in the anchor and positive samples. The channel values for each biased color are listed in Table I. The unbiased one is the case where we color randomly every image, which means that color is the class-independent feature. Note that the anchor and positive samples share the digit information, not style like rotation, thickness, and font.

As shown in Figure 6, in the biased color case (first row), the color appears in the content part because the color is a shared feature between the anchor and positive images. We can see that color changes only when the class (digit) changes in the latent space, and also style features appear in the style part. In the unbiased color case (second row), the color appears in the style part because color is one of the prevalent variations in the dataset. As a result, the color turns up in the style part with other style features. We observe that by partitioning the representation into the two parts, our approach successfully separates two types of features in the VAE framework.

As an application of the partitioned representation, it is possible to generate a new sample by switching the content and style parts of two samples as in Figure 7. In Figure 7(a), the first two images of each low are the input images to VAE trained on unbiased color MNIST, and the last two images are the reconstructed ones with a switched style. Figure 7(b) shows how to obtain the new images. We exchange the style part of the mean vectors from the two input images. Then, we reconstruct the new samples from the new mean vector which consists of the content part of one image and the style part of the other image. As a result, the new samples represent the same digit with the style of the other image.

In this section, we investigate the effect of the partitioned representation in classification tasks. We first train an additional linear classifier over the mean vector of the trained VAE model to classify the mean vector. To see the effect of the model, we add Gaussian noise at the style part or the content part of the latent variable. The test cases can be summarized as follows:

• •

  Noise on Style: Gaussian noise is added to the style part of the mean vector.

• •

  Noise on Content: Gaussian noise is added to the content part of the mean vector.

By adding noise, we perturb the style or content part of the mean vector to check if the proposed method can split content and style features.

Table II shows the results with both datasets. For the VAE model trained on the biased color dataset, the classifier achieves 99% on the test set from the biased color dataset. When Gaussian noise is added, we observe that the style part is more robust to the noise than the content part, because the style features have nothing to do with the classification task while content features are class-related. Similarly, for the unbiased color dataset, adding noise to the style part is more robust than adding noise to the content part. Another interesting point is that the biased color case degrades less than the unbiased color case when the noise is added to the content part. Since the biased color case has a strong class-related feature with color in the content part, the representation can be more robust against random noise on it.

Additionally, Table III shows how accuracy changes as the intensity of the noise increases. Given the Gaussian noise $n\sim N(0,1)$ and the level of intensity $t\in[1,4]$, we add $t\times n$ to the content or style part. For both color datasets, the style part is relatively more robust than the content part against the noise especially when the noise intensity increases.

In this section, we evaluate the generalization ability of the partitioned representation in various tasks. In all experiments, we compare our approach (i.e., BYOL with the partitioned representation) to the conventional BYOL in terms of the accuracy of downstream tasks. For the partitioned representation, we divide 256 dimensions of representation of the conventional BYOL: 192 dimensions for the content part and 64 dimensions for the style part. Our implementation is based on the official implementation given by [8].

First, we train the two BYOL models based on the Resnet18 encoder on the STL10 dataset, which consists of 100k unlabeled images and 13k labeled images for ten classes (5k for training and 8k for the test, respectively) [24]. We follow the conventional BYOL for the implementation of data augmentation and training strategy [8]. We also follow the linear evaluation procedures described in [7] to measure the linear separability of representations. We first pretrain both BYOLs on unlabeled data over 800 epochs and then train a linear classifier on 5k training data over a frozen encoder. Finally, we measure the classification performance of the test data for the linear evaluation task. As a result, our method achieves the accuracy of $79.7\%$, which is $1.4\%p$ higher than the conventional BYOL.

We also conducted the same experiment on the BYOL models over 40 epochs with 256 batch sizes on ImageNet [25], which consists of 1.28M training images and 50k validation images. As presented in Table IV for the linear evaluation task, our method gets $1.0\%p$ greater accuracy than the conventional BYOL. We think that a model becomes flexible to extract more effective features with the partitioned representation which leads to better generalization.

Lastly, we experiment a few-shot learning task as a downstream task of a pretrained encoder to understand the transferability of the partitioned representation, because few-shot learning from the pretrained encoder can measure the generalization ability of the pretrained encoder [26]. Few-shot learning combined with meta-learning classifies images with a few training samples, and the test images are from unseen classes during training. It typically performs ‘$K$-shot $N$-way’ tasks where we randomly take $N$ classes among entire classes and draw $K$ samples for each class. In every iteration, we train an encoder with $N\times K$ samples, and then test the model on unseen classes.

It was proposed to pretrain an embedding network, AMDIM [2], in a self-supervised way which is supposed to be fine-tuned for a few-shot classification task [26]. It has two training phases: self-supervised pretraining and meta fine-tuning. We replaced the embedding network of [26] with the conventional BYOL or our modified BYOL.

We use the MiniImagetNet dataset [27], which contains 60k images from 100 classes. Also, we follow the data composition as the baseline to divide 100 classes into 64 classes for training, 16 classes for validation, 20 classes for the test [26]. First, an embedding network is pretrained on 100 classes, and meta-learning is applied to fine-tune the network as described in [26]. We follow the same training process on the conventional BYOL and BYOL based on Resnet 50 with the partitioned representation on ‘1-shot 5-way’ and ‘5-shot 5-way’ classification tasks.

As shown in Table V, the conventional BYOL outperforms the baseline on both tasks because the conventional BYOL has better transferability than AMDIM. We find that our method improves further. The partitioned representation extracts precise class-related features and class-unrelated features at the same time. That is, the partitioned representation disentangles the information for classification, which makes classification more efficient.