ProFeAT: Projected Feature Adversarial Training for Self-Supervised Learning of Robust Representations

TL;DR

ProFeAT introduces a novel training framework that enhances self-supervised adversarial learning by aligning teacher-student objectives with a projection head, significantly improving robustness and accuracy over previous SSL-AT methods.

Contribution

The paper proposes Projected Feature Adversarial Training (ProFeAT), a new method that reduces the performance gap in SSL-AT by using a projection head and tailored losses, achieving state-of-the-art results.

Findings

Significant improvements in clean and robust accuracy.

ProFeAT outperforms existing SSL-AT methods.

Comparable or better performance than supervised-AT methods.

Abstract

The need for abundant labelled data in supervised Adversarial Training (AT) has prompted the use of Self-Supervised Learning (SSL) techniques with AT. However, the direct application of existing SSL methods to adversarial training has been sub-optimal due to the increased training complexity of combining SSL with AT. A recent approach, DeACL, mitigates this by utilizing supervision from a standard SSL teacher in a distillation setting, to mimic supervised AT. However, we find that there is still a large performance gap when compared to supervised adversarial training, specifically on larger models. In this work, investigate the key reason for this gap and propose Projected Feature Adversarial Training (ProFeAT) to bridge the same. We show that the sub-optimal distillation performance is a result of mismatch in training objectives of the teacher and student, and propose to use a projection head at the student, that allows it to leverage weak supervision from the teacher while also being able to learn adversarially robust representations that are distinct from the teacher. We further propose appropriate attack and defense losses at the feature and projector, alongside a combination of weak and strong augmentations for the teacher and student respectively, to improve the training data diversity without increasing the training complexity. Through extensive experiments on several benchmark datasets and models, we demonstrate significant improvements in both clean and robust accuracy when compared to existing SSL-AT methods, setting a new state-of-the-art. We further report on-par/ improved performance when compared to TRADES, a popular supervised-AT method.