Improve Model Robustness in Less Time Than It Takes to Drink A Cup of Coffee with Plug-and-Play Robustness Plugins

Abstract

Self-supervised learning has become the primary method for pre-training large models due to its ability to train without labeled data and its excellent data feature representation capabilities. However, neural network models are vulnerable to adversarial attacks, which can lead to incorrect predictions. Previous work has attempted to enhance the robust representation capabilities of base models through self-supervised adversarial training (self-AT), which integrates adversarial training into the self-supervised learning pre-training process. However, self-supervised learning requires numerous training epoches, and adversarial training is computationally complex. Consequently, these methods need an additional 2.75 to 12 times the pre-training time of the model to obtain robust representations. Considering the resource consumption of training large models and the current high cost of computational resources, the cost of obtaining robustness for base models is excessively high and impractical. This paper proposes a novel Plug-and-Play model Robustness Plugin training framework called PPRP. PPRP is designed as a robustness plugin for self-supervised base models that have completed pre-training. Once the robust plugin is added, the base model gains robust representation capabilities. Essentially, PPRP is a teacher-student network that performs adversarial training on a plugin model with only a few parameters, reducing the time required to achieve model robustness to 5% of the pre-training time. The robust plugin can be seamlessly integrated into pre-trained models without additional inference latency. Experiments show that on multiple datasets, different base models with the PPRP-trained robust plugin achieve state-of-the-art robustness.