Transferable Unsupervised Robust Representation Learning

1 INTRODUCTION . Recently , there has been an increasing interest in unsupervised visual representation learning , where the goal is to learn effective representations of images without using human annotations ( Bachman et al. , 2019 ; Chen et al. , 2020b ; He et al. , 2020 ; Misra & Maaten , 2020 ) . In this work , we consider a relatively under-explored aspect : the robustness of these representations . Here , we use robustness to encompass the representation ’ s resilience to common corruptions ( Hendrycks & Dietterich , 2019 ) , adversarial examples ( Madry et al. , 2018 ) , and its ability to measure the uncertainty of its output in the face of such perturbations ( Kumar et al. , 2019 ) . More specifically , we develop new learning methods for transferable robustness of these representations , which improves and preserves the representation ’ s robustness after fine-tuning on downstream tasks . While recent works have only evaluated the virtue of a representation by its accuracy on clean data after transfer learning and fine-tuning , we argue that transferable robustness should also be an integral part of a good representation . Recently , Rezaei & Liu ( 2020 ) have shown that downstream task models are vulnerable to adversarial attacks that are based solely on the pre-trained representation before fine-tuning . This highlights the importance of transferable robustness , without which downstream task models are vulnerable and lack generalization . Moreover , it is commonly assumed that there needs to be a trade-off between natural accuracy ( on clean data ) and robust accuracy ( on corrupted data ) ( Zhang et al. , 2019 ) . While there have been works showing that such a trade-off is not necessary for supervised learning ( Yang et al. , 2020 ) , it is unclear how the principles they leverage is applicable to unsupervised representation learning . Summary of Contributions : ( 1 ) We show that the accuracy-robustness trade-off is not necessary for unsupervised representation learning . ( 2 ) We develop a new representation learning framework : Unsupervised Robust Representation Learning ( URRL ) , which uses a novel mix of representation learning pretext task and robust supervised learning ( AugMix , Hendrycks et al . ( 2020b ) ) . ( 3 ) In addition , we propose a task-agnostic similarity regularization that further improves the robustness of downstream tasks without the need for any adversarial training or additional data augmentation . ( 4 ) We introduce an evaluation framework for the transfer robustness of a representation , which includes its resilience to corruptions , adversarial robustness , and uncertainty calibration . ( 5 ) We show that URRL improves both the clean accuracy and robustness of state-of-the-art representation learning under linear evaluation , and our full framework , URRL with Similarity Regularization ( URRL-SR ) , further improves its robust accuracy by over 10 % on 12 datasets . The proposed URRL is shown in Figure 1 ( a ) . Our approach is motivated by the fact that in previous works , data augmentation has played an important role in both supervised learning robustness and unsupervised representation learning ( Yun et al. , 2019 ; Chen et al. , 2020b ) . URRL randomly samples a pretext task of unsupervised representation learning and a robust data augmentation ( AugMix , Hendrycks et al . ( 2020b ) ) , and then optimizes them both under a contrastive learning framework . We show that URRL improves both the clean accuracy ( on the original images without corruptions ) and the robustness of the learned representation under the standard linear evaluation protocol . Nevertheless , using the representation under linear evaluation limits the performance of transfer learning , and ultimately the usefulness of representation learning in downstream tasks . It is thus important to evaluate the representations ’ transferable robustness after fine-tuning the entire network to the downstream task . The biggest challenge is to preserve the inherited robustness of the representations : an aggressive fine-tuning schedule on a small dataset can significantly overwrite the robustness of the learned representations . We address this challenge by proposing a task-agnostic similarity regularization that is applicable for the fine-tuning of the learned representation on any downstream task . We show that this improves the robustness of the fine-tuned model . The proposed regularization is shown in Figure 1 ( b ) . fθ is initialized by the robust representation f∗ learned in Figure 1 ( a ) , and is now being fine-tuned by the downstream task loss Ltask . f∗ is robust because it is trained on augmentation of corrupted versions of an image . We hypothesize that this form of robustness is not just about how f∗ represents an image individually , but also about how f∗ measures the similarities between images . This implies that we can preserve the robustness of f∗ by maintaining its similarity measures between images . We achieve this by regularizing the similarity matrix of a representation during fine-tuning to be similar to that of the original representation before fine-tuning . This allows the downstream task model to better utilize the robustness of the learned representation , while not capping the performance , which would be the case if the representation is fixed . Our full approach , URRL with similarity regularization ( URRL-SR ) , further improves URRL ’ s robust accuracy by over 10 % on 12 datasets without any adversarial training or further data augmentation . 2 METHOD . In this work , we advocate the transferable robustness of unsupervised representation learning , which measures the robustness of a representation after fine-tuning on downstream tasks . The two challenges are : ( i ) It is commonly assumed that there is an accuracy-robustness trade-off , and the promotion of transferable robustness could potentially set back the performance gain from unsupervised representation learning . ( ii ) the robustness of a representation could be changed drastically during the fine-tuning process , and thus make it challenging to preserve the robustness after fine-tuning the whole network . In Section 2.1 , we address the first challenge by proposing URRL , which combines the advantages of pretext tasks from unsupervised representation learning with robust augmentations to improve not only the robustness but also the clean accuracy . In Section 2.2 , we address the second challenge by proposing a task-agnostic similarity regularization that preserves the robustness of learned representations , while not capping the performance of transfer learning . 2.1 UNSUPERVISED ROBUST REPRESENTATION LEARNING . In this section , we discuss the proposed Unsupervised Robust Representation Learning ( URRL ) , which combines the advantages of representation learning pretext task and robust augmentation . We observe that data augmentation plays a central role in the recent literature of both unsupervised representation learning and supervised learning robustness . Contrastive learning has become a predominant design for representation learning ( Chen et al. , 2020b ) . Given an unlabeled dataset { xi } of images xi , the goal is to learn the representation of each image zi = fθ1 ( xi ) by optimizing the parameters θ1 of the encoding function . The role of data augmentation is to randomly create two correlated views of the same image . Let Trep be a family of data augmentation . We can then sample two data augmentation functions t1 , t2 ∼ Trep , and apply them to an image xi to obtain a positive pair xq = t1 ( xi ) and xk+ = t2 ( xi ) . The goal is then to train the encoder fθ1 such that it still encodes the two augmented images xq and xk+ similarly despite the random augmentations to the images . In this case , the representation can learn to ignore information that is not task-relevant ( introduced by the random augmentation ) , while still capturing the task-relevant information in the images . One can achieve this through the contrastive loss : Lcontra = − log exp ( zq · zk+/τ ) exp ( zq · zk+/τ ) + ∑ k− exp ( zq · zk−/τ ) , ( 1 ) where τ is the softmax temperature , and zk− are the representations of negative samples . The negative samples are often randomly augmented versions of images in the same batch xj , j 6= i . The family of data augmentation Trep plays a crucial role here as it would determine what the encoder fθ1 learns to ignore and preserve in the representation of an image . On the other hand , data augmentation improves the robustness of supervised learning by effectively expanding the training dataset without human annotation . Given a dataset { xi , yi } of images xi and the corresponding labels yi , one can apply data augmentation Trob to obtain an augmented dataset { Trob ( xi ) , yi } , where Trob ( xi ) = t ( xi ) , t ∼ Trob . Trob improves the robustness of the model by indicating that the augmented image Trob ( xi ) still has the same label yi as the original image xi . Recent works have shown that by mixing various types of perturbations , the learned classifier has the potential to be robust to novel perturbations that are not used in training ( Hendrycks et al. , 2020b ) . Trob plays a crucial role in which types of perturbation the learned model would be robust to . Given the importance of data augmentation in both of these aspects , a reasonable direction for unsupervised robust representation learning is to design a new family of data augmentation functions that can be used in Eq . ( 1 ) to learn robust representation without human annotation . We observe that the goal of Trep and Trob are quite similar at the high-level : In robustness literature , the aim of data augmentation is to increase the diversity of training data as much as possible without drifting off the data manifold ( Hendrycks et al. , 2020b ) . On the other hand , recent work hypothesizes that the best data augmentations for unsupervised representation learning are ones that “ share the minimal information necessary to perform well at the downstream task ” ( Tian et al. , 2020 ) . If we take the downstream task to be classification and assume that the reason for not drifting off the data manifold is also optimizing for classification performance , then the aims of these two augmentations are actually quite similar : both are trying to preserve information that is enough to do well on classification , while increasing perturbation ( and decreasing shared information ) as much as possible . As we will show in the experiments , this is indeed the case : a representation learned with Trep is already more robust than ones learned by supervised pre-training ; At the same time , Trob can already serve as an effective pretext task for unsupervised representation learning . It could seem that using Trob as a pretext would directly lead to unsupervised representation learning with transferable robustness . However , one missing piece of naively using Trob as a pretext for representation learning is the augmentation that focuses on color distortion . As observed and shown by Chen et al . ( 2020b ) , color distortion plays a critical role in the success of unsupervised representation learning because patches from the same image could share a similar color distribution , and without color distortion augmentation , the encoder could learn to exploit a shortcut to minimize Eq . ( 1 ) by focusing on the color histogram of an image . Drastic color distortion is often discouraged for Trob because “ histogram color augmentation can change the image ’ s class ” ( Hendrycks et al. , 2020b ) . Therefore , we propose a straightforward solution to the problem by incorporating color histogram distortion into Trob . The proposed data augmentation family Turrl for URRL does not sample the data augmentation function t directly . Instead , we use a two-stage process : Turrl ( xi ) = t ( xi ) , t ∼ T , T = { Trob with probability p , Trep with probability 1− p , ( 2 ) where we first sample the family of data augmentation T , and then sample the actual augmentation from T . Here , the sampled T could be either Trob or Trep . With the inclusion of Trep , Turrl prevents the encoder f from exploiting the color histogram . On the other hand , the use of Trob in addition to Trep further strengthens the robustness of the learned representation . We will show that this simple approach is able to outperform more sophisticated approaches to modify Trob . In this work , we use the pretext task of Chen et al . ( 2020d ) as Trep and AugMix ( Hendrycks et al. , 2020b ) as Trob

Summary: This paper proposes *Unsupervised Robust Representation Learning* (URRL), a framework that combines several data augmentation schemes and a similarity-based loss. The goal is to improve the robustness of visual representations to image perturbations. A further goal is to maintain the robustness properties of pre-trained representations after fine-tuning the network to downstream tasks.