Supervised Contrastive Prototype Learning:
Augmentation Free Robust Neural Network

DNNs are known for their poor performance in generalizing to unseen inputs or simple image transformations, such as rotation (Nguyen, Yosinski, and Clune 2015; Guo et al. 2017). Two inputs with insignificant $l_{p}$ distance in the input space can produce significantly distant feature representations that lead to misclassification. On the contrary, input that is perceptually different from the training samples, such as random noise, can result in feature representations nearly identical to the training samples. In this work, we use the term robustness to refer to the property of DNN to learn an equidistant projection from the input space to the feature space. We identify two main failure cases in DNN; lack of robustness against perturbations, and against nuances.

Robustness against perturbations can be evaluated with an input crafted by an adversarial attack, where an insignificant perturbation on the input space can cause a significant perturbation on the feature space (Szegedy et al. 2013). Similarly, robustness against nuances can be observed with the weakness of DNN where they assign nearly identical feature representations to input that is significantly distant in the $l_{p}$ input space, such as between random noise and training samples.

Previous work improves the robustness of DNNs with training augmentation techniques specific to the problem they solve. For example, some work (Papernot et al. 2016; Madry et al. 2017; Papernot et al. 2017) improves adversarial robustness by including adversarial-generated samples in the training set. OOD detection methods use samples from the OOD set (Lee et al. 2017; Liang, Li, and Srikant 2017) or an auxiliary dataset (Hendrycks, Mazeika, and Dietterich 2018) during training.

Previous work in contrastive learning (Chen et al. 2020b) identifies training augmentation techniques as one of the primary contributing factors to learning robust feature representations invariant to nuances in the input. In this work, we also view that training augmentation techniques are necessary for learning representations invariant to noise. However, we identify two issues with current approaches in training augmentation.

First, DNNs are prone to over-fitting. We find that, for multiple architectures, the majority of OOD samples are assigned to the most difficult class, the one with the least salient features. Our observation is at odds with the goal of a classifier to learn salient features for every class. For example, for a pretrained ResNet classifier, 99% of uniform random noise is assigned to class 104 with high probability (Appendix fig. S2 c). This suggests a failure of learning, whereby all kinds of unknown images are assigned to the most uncertain class (here, 104). Additionally, DNNs over-fit the augmentation technique used to train the classifier. Previous work (He et al. 2016; Simonyan and Zisserman 2014) resize training images to 256 pixels on the shortest side and then randomly crop them to (224x224) pixels when training on ImageNet. fig. 5 visualizes artifacts of over-fitting to the augmentation technique used during training. The model is biased to the least salient features when predicting out-of-distribution images.

Secondly, data augmentation techniques are expensive, especially for OOD methods that require samples from an out-of-distribution set or an auxiliary dataset. Such methods can be limited in practical settings where additional data is scarce. Similarly, generating adversarial examples for every training iteration is expensive and the current state-of-the-art (Wong, Rice, and Kolter 2020) focuses on improving computational performance. Current training augmentation solutions are a quick fix to a systematic problem, where the issue of robustness is an issue of generalization (Xu and Mannor 2012).

Contrastive Learning (Hadsell, Chopra, and LeCun 2006) approaches do not directly solve for a task but optimize a similarity metric such as the euclidean distance or cosine similarity between feature vectors. Previous contrastive learning work has been successful on downstream tasks such as classification, zero-shot learning, and more (Wang et al. 2018; Xu et al. 2020; Wang et al. 2017; Chen et al. 2017). Additionally, the metric space learned by such models can be robust to both adversarial examples and OOD. The focus of state-of-the-art contrastive learning approaches (Chen et al. 2020a; Khosla et al. 2020) has been to find effective ways to sample pairs that are of the same class (“positive pair”) and one of the most-similar other classes (“negative pair”) that produce meaningful gradients. The sample mining process is computationally expensive and the downstream task performance is sensitive to both the sampling method and the sample size (Kaya and Bilge 2019)

Prototypical Learning (PL) uses a set of vectors (prototypes) to represent a class and is an extension of k-means, where each of the K vectors is a prototype. PL has been used in combination with different learning paradigms such as pattern recognition (Liu and Nakagawa 2001) and more recently has been used both for zero-shot and few-shot learning due to the generalization and sample efficiency (Snell, Swersky, and Zemel 2017). Frameworks have been proposed that implement PL in both supervised and unsupervised settings (Yang et al. 2018a; Li et al. 2020).

Our goal is to design a model that does not depend on any training augmentation technique or auxiliary datasets to achieve high robustness. Motivated by the recent success of Contrastive Learning and Prototypical Learning approaches, in this work we propose the Prototype Classification Head (PCH). Our proposed method combines a classification task with a purely contrastive, sample efficient loss, and can be trained end-to-end with a DNN backbone. PCH is used at the penultimate feature representation of a DNN and replaces the classification head. We use a set of prototype vectors to represent a cluster of feature vectors of the same class. We further improve on previous prototype learning approaches and use N-Prototype Loss which increases inter-class separability and intra-class compactness among training samples. Our approach creates more robust models without additional tricks. Our contribution in detail:

• •

  To the best of our knowledge, this is the first work to combine prototype learning with a purely contrastive loss between prototypes and samples, N-Prototype Loss

• •

  Compared to traditional prototype and contrastive learning methods, our framework uses prototypes to represent a class and does not require intra-batch sample mining.

• •

  Our proposed framework is flexible and can be added to any existing DNN architectures without any further modifications.

• •

  We provide a theoretical justification for the improved robustness, and conduct quantitative experiments that outperform other state-of-art methods

We identify Prototype Learning approaches and Contrastive Learning approaches as two closely related categories of work that improve the robustness of DNN classifiers.

Triplet loss (Hermans, Beyer, and Leibe 2017) uses one negative and positive pair for each sample to learn a metric and can be sensitive to noise which leads to over-fitting. N-pair(Sohn 2016a) loss decreases the similarity between multiple negative pairs and increases the similarity between a positive pair. Current contrastive learning approaches suffer from poor quality sample pairs that generate uninformative sub-gradients which lead to poor model convergence. Mining for informative pairs such as hard negatives can be computationally expensive. Several approaches have been proposed to improve on the sample mining problem. (Duan et al. 2018) train a separate negative sample generator jointly with a DNN that augments “easy” negative samples with “hard” negative samples. Lifted Structured Loss (LSL) (Oh Song et al. 2016) uses multiple positive and negative pairs in the importance sampling method to improve sample mining of hard negative pairs. However, sample mining techniques are computationally expensive. In contrast, we use learned prototypes for positive and negative pairs and thus avoid the sample mining problem.

Contrastive learning losses are applied in a supervised setting. (Mao et al. 2019) use Triplet loss in an adversarial training setting as an auxilary training objective to the cross-entropy loss. Several contrastive learning framework such as SimCLR (Chen et al. 2020b), MoCo (He et al. 2020) and BYOL (Grill et al. 2020) use a generalized version of N-pair loss (InfoNCE) on the feature representation of an augmented sample to create a positive pair and treat all other samples in a batch as negative. Work by (Li et al. 2020) add an additional loss term to the NCE loss that compares the original sample with prototypes of the same cluster and r negative prototypes. Both the positive and negative prototypes are determined using an expectation maximization clustering framework at every training iteration. In contrast, we assign prototypes to each corresponding class in a supervised setting. Additionally, we use a single training objective between each sample and multiple prototypes, and have a flexible penultimate feature dimension which we identify to be the root cause of the improved robustness. Lastly, our method is orthogonal to other unsupervised contrastive learning approaches and any training augmentation techniques or specific backbone architectures (Chen et al. 2020b; He et al. 2020; Grill et al. 2020) can be extended to our framework.

(Mustafa et al. 2019) apply Prototype Conformity Loss (PCL) as the cross-entropy of the cosine similarity between the feature representation of a sample and the true class prototype vector. For every intermediate representation, an auxiliary network is trained to enforce the intra-class compactness and inter-class separation using the learned prototypes. In contrast to our method, it requires modifying the existing network architecture and can not be directly applied to any existing architecture. Additional work (Mettes, van der Pol, and Snoek 2019; Chen et al. 2021) define fixed prototypes based on an inductive bias for the latent space geometry and maximize the cosine similarity between a sample and the positive class prototype. In contrast to our method, the loss function is applied between multiple sample feature representations and with multiple prototypes from the positive and negative class. Lastly, in our framework the prototypes are learned end-to-end with the DNN. Maintaining fixed prototypes is orthogonal to our method and can be considered as a special case of our framework.

The work that is most similar to ours is (Yang et al. 2018b) that use Distance-based Cross-Entropy (DCE) loss and (Liu et al. 2020) that extend on DCE with an additional loss term Prototype Encoding Cost (PEC). DCE is applied between the closest positive class prototype and a sample, while PEC is the $l_{2}$ norm between the positive class prototype and the penultimate feature representation. DCE and PEC are optimized jointly for both the network parameters and learned prototypes. Similarly, our framework is applied to the penultimate feature representation. Contrary to DCE and PEC, at every step our method increases the distance between a sample and multiple ($N$) opposite class prototypes and reduces the distance to the closest positive prototype.

There are additional work (Yang et al. 2020; Shu et al. 2020) that use prototype learning for a specific problem but can be seen as derivation of more general frameworks. In this work and for reasons of clarity we provide a generalized definition of the state-of-the-art approaches in section 3. Extensive experiments show that our method outperforms other state-of-the-art (Mao et al. 2019; Chen et al. 2020b; Yang et al. 2018a) methods in clean robustness and does not require training augmentation techniques to achieve high robustness. Finally, when compared to other metric learning approaches, our approach does not require sample mining.

In this work, we consider the issue of robustness from two angles. Robustness against data samples crafted by an adversary and robustness against samples that are out-of-distribution compared to the training data set.

Out-of-Distribution Robustness (OOD) evaluates a classifier $f_{\theta}(x)=h_{x}$ trained on a dataset $\mathcal{D}^{\emph{in}}$ with an auxiliary function $g(h_{x})$ on a test dataset $\mathcal{D}^{\emph{out}}$ with $\mathcal{D}^{\emph{out}}\cap\mathcal{D}^{\emph{in}}=\emptyset$ such that

In this work, we do not design a specific method to our model for OOD detection and use work by (Hendrycks and Gimpel 2018).
OOD Baseline (Hendrycks and Gimpel 2018) normalize the predicted class similarity for a sample $x$ and classifier $f_{\theta}(x)$ and use the maximum value of the prediction scores for all classes $K$ as the confidence score $p(h_{x})=\max{\sigma(h_{x})}$

where a confidence threshold $\alpha$ is used to determine if $x\in D^{\emph{out}}$. The robustness of a classifier is then determined by the performance of $g(f_{\theta}(x))$ in correctly classifying in and out-of-distribution samples for a fixed $\alpha$.

Adversarial Robustness measures the accuracy of a DNN with $f_{\theta}$ in correctly predicting a clean sample $x$, after a perturbation $\delta$ has been applied by an adversary, with perturbation budget $\varepsilon$ such that

where $\mathcal{L}(\cdotp)$ denotes the loss function used to train $f_{\theta}$ and $\parallel\delta\parallel_{\textit{p}}$ is the p-norm of $\delta$ with $p\geq 1$. As such, the adversarial example lies within an $\ell_{p}$-ball centered at $x$.

A robust model can correctly classify a perturbed sample that is within the perturbation budget. In this work, we consider adversaries that can manipulate a data sample in $l_{\infty}$ and $l_{2}$ space.
We consider the following adversarial attack settings: Fast Gradient Sign Method (FGSM) (Goodfellow, Shlens, and Szegedy 2014), Basic Iterative Method (BIM) (Kurakin, Goodfellow, and Bengio 2016) , Projected Gradient Descent (PGD) (Madry et al. 2017) and AutoAttack (Croce and Hein 2020).

Supervised Constrastive Prototype Learning (SCPL) uses a DNN $f_{\theta}(x)=h_{x}$ as feature extractor, where $x$ is the raw input, $\theta$ the parameters of the model and $h_{x}$ the learned feature representation of $x$. We apply a Prototype Classification Head on the hidden feature representation $h_{x}$ of $x$.

The Prototype Classification Head (PCH) is composed of multiple learnable prototype vectors $\mathcal{M}=\{m_{i}^{j}\in R^{D}|i\in K;j\in C_{i}\}$ where $i$ is the index from the set of classes ${K}$, $j$ is the index from the set of prototypes for class $i$ and $C_{i}$ is the total number of prototypes of the same class. For our experiments, we use the same number of prototypes $C$ for all classes, but this assumption can be relaxed for future works.

Our proposed method outperforms the previous state-of-the-art, but there are limitations and potential future improvements. First, the prototypes are initialized with random weights in the current setting, which can lead to poor convergence at the beginning of the training stage. Warm-up methods could further improve model convergence (Arthur and Vassilvitskii 2006).

In addition, the similarity function between feature representations and the decision boundary $\alpha$ from theorem Provably Robust Classification in the Appendix can be improved. Related to this is the design of an OOD method specific to our model.

Finally, the inference of our proposed method is based on prototype matching by assigning test examples to the class of their closest prototype, which cannot be directly applied to tasks where a continuous value needs to be predicted (i.e regressional analysis).

We propose a supervised contrastive prototype learning (SCPL) framework for learning robust feature representation. We replace the categorical classification layer of a DNN with a Prototype Classification Head (PCH). PCH is an ensemble of prototypes assigned to each class. We train PCH jointly with the DNN using a single optimization objective, N-Prototype Loss. We empirically validate our results and provide a theoretical justification for the improved robustness. We achieve increased robustness without requiring auxiliary data such as adversarial or out-of-distribution examples. Lastly, SCPL requires no modification to the classification backbone and thus is compatible with any existing DNN.