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 uses a different data augmentation (AugMix) scheme to improve self-supervised representation learning. It improves accuracy and (corruption and adversarial) robustness by a sufficiently interesting amount. The paper's presentation is clear, but the paper could be more thorough. Since the technique is simple and general, it could easily be broadly applicable to the burgeoning area of self-supervised learning.