Augmentation-aware Self-supervised Learning with Conditioned Projector

TL;DR

This paper proposes to foster sensitivity to characteristics in the representation space by modifying the projector network, a common component of self-supervised architectures, and conducts a series of experiments, which show that CASSLE improves over various SSL methods, reaching state-of-the-art performance in multiple downstream tasks.

Abstract

Self-supervised learning (SSL) is a powerful technique for learning from unlabeled data. By learning to remain invariant to applied data augmentations, methods such as SimCLR and MoCo can reach quality on par with supervised approaches. However, this invariance may be detrimental for solving downstream tasks that depend on traits affected by augmentations used during pretraining, such as color. In this paper, we propose to foster sensitivity to such characteristics in the representation space by modifying the projector network, a common component of self-supervised architectures. Specifically, we supplement the projector with information about augmentations applied to images. For the projector to take advantage of this auxiliary conditioning when solving the SSL task, the feature extractor learns to preserve the augmentation information in its representations. Our approach, coined Conditional Augmentation-aware Self-supervised Learning (CASSLE), is directly applicable to typical joint-embedding SSL methods regardless of their objective functions. Moreover, it does not require major changes in the network architecture or prior knowledge of downstream tasks. In addition to an analysis of sensitivity towards different data augmentations, we conduct a series of experiments, which show that CASSLE improves over various SSL methods, reaching state-of-the-art performance in multiple downstream tasks.

Figures (10)

• Figure 1: In the traditional self-supervised setting, contrastive loss minimization pulls the representations of augmented image views closer in the latent space of the projector (left). This may also reduce the distance between their feature extractor representations (right). Thus, the representation becomes invariant to augmentation-induced perturbations, which may hinder the performance on downstream tasks. In contrast, the self-supervised objective of CASSLE draws together joint representations of images and their augmentations in the projector space (bottom row). By conditioning the projector with augmentation information, image representations retain more sensitivity to perturbations in the feature extractor space. This proves to be beneficial when solving downstream tasks.
• Figure 2: Overview of CASSLE. We extend the typical self-supervised learning approaches by incorporating the information of augmentations applied to images into the projector network. In CASSLE, the SSL objective is thus imposed on joint representations of images and the augmentations that had been applied to them. This way, CASSLE enables the feature extractor to be more aware of augmentations than the methods that do not condition the projector network.
• Figure 3: Cosine similarities of CASSLE projector $(\pi)$ representations when conditioned with augmentation information from either their respective images (green) or randomly sampled (red). Solid lines denote the mean values of similarities. Conditioning the CASSLE projector with wrong augmentation information decreases its ability to draw image pairs together, indicating that it indeed relies on augmentation information to perform its task.
• Figure 4: A visualization of four considered methods of conditioning the projector with augmentation information.
• Figure 5: Image retrieval examples for Cars and Flowers datasets.