Supervised Contrastive Prototype Learning: Augmentation Free Robust Neural Network

TL;DR

This paper introduces Supervised Contrastive Prototype Learning (SCPL), a novel training framework that enhances neural network robustness against adversarial and out-of-distribution samples without requiring data augmentation or architectural changes.

Contribution

The paper proposes SCPL, a sample-efficient contrastive learning method with prototype classification, improving robustness and nuance invariance in neural networks compared to existing approaches.

Findings

SCPL outperforms state-of-the-art contrastive methods in robustness.

It is sample efficient and does not require sample mining.

Can be integrated with existing models and training techniques.

Abstract

Transformations in the input space of Deep Neural Networks (DNN) lead to unintended changes in the feature space. Almost perceptually identical inputs, such as adversarial examples, can have significantly distant feature representations. On the contrary, Out-of-Distribution (OOD) samples can have highly similar feature representations to training set samples. Our theoretical analysis for DNNs trained with a categorical classification head suggests that the inflexible logit space restricted by the classification problem size is one of the root causes for the lack of $\textit{robustness}$. Our second observation is that DNNs over-fit to the training augmentation technique and do not learn $\textit{nuance invariant}$ representations. Inspired by the recent success of prototypical and contrastive learning frameworks for both improving robustness and learning nuance invariant representations, we propose a training framework, $\textbf{Supervised Contrastive Prototype Learning}$ (SCPL). We use N-pair contrastive loss with prototypes of the same and opposite classes and replace a categorical classification head with a $\textbf{Prototype Classification Head}$ (PCH). Our approach is $\textit{sample efficient}$, does not require $\textit{sample mining}$, can be implemented on any existing DNN without modification to their architecture, and combined with other training augmentation techniques. We empirically evaluate the $\textbf{clean}$ robustness of our method on out-of-distribution and adversarial samples. Our framework outperforms other state-of-the-art contrastive and prototype learning approaches in $\textit{robustness}$.